[test bugs] uda_mvtec (PaddlePaddle#3133)

PeiyuLau · Apr 7, 2023 · 80f9551 · 80f9551
1 parent 3a94093
commit 80f9551
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/contrib/QualityInspector/qinspector/uad/datasets/mvtec.py b/contrib/QualityInspector/qinspector/uad/datasets/mvtec.py
@@ -39,7 +39,8 @@ def __init__(self,
                  class_name='bottle',
                  is_train=True,
                  resize=[256, 256],
-                 cropsize=[224, 224]):
+                 cropsize=[224, 224],
+                 is_predict=False):
         assert class_name in CLASS_NAMES, 'class_name: {}, should be in {}'.format(
             class_name, CLASS_NAMES)
         self.dataset_root_path = dataset_root_path
@@ -49,7 +50,7 @@ def __init__(self,
         self.cropsize = cropsize
 
         # load dataset
-        if is_train:
+        if not is_predict:
             self.x, self.y, self.mask = self.load_dataset_folder()
 
         # set transforms

diff --git a/contrib/QualityInspector/tools/uad/padim/predict.py b/contrib/QualityInspector/tools/uad/padim/predict.py
@@ -89,7 +89,7 @@ def main():
     print("Testing model for {} with sigle picture".format(class_name))
 
     # build datasets
-    MVTecDataset = mvtec.MVTecDataset(is_train=False)
+    MVTecDataset = mvtec.MVTecDataset(is_predict=True)
     transform_x = MVTecDataset.get_transform_x()
     x = Image.open(args.img_path).convert('RGB')
     x = transform_x(x).unsqueeze(0)

diff --git a/contrib/QualityInspector/tools/uad/patchcore/predict.py b/contrib/QualityInspector/tools/uad/patchcore/predict.py
@@ -97,7 +97,7 @@ def main():
     model.eval()
 
     # build data
-    MVTecDataset = mvtec.MVTecDataset(is_train=False)
+    MVTecDataset = mvtec.MVTecDataset(is_predict=True)
     transform_x = MVTecDataset.get_transform_x()
     x = Image.open(args.img_path).convert('RGB')
     x = transform_x(x).unsqueeze(0)