adding args of datapath

Author: NarcissusEx

README.md

@@ -52,11 +52,11 @@ bash setup.sh
 ### 2) Template Generation
 **Original template set on MVTec AD:**
 ```bash
-python run.py --mode temp --ttype ALL --dataset MVTec_AD
+python run.py --mode temp --ttype ALL --dataset MVTec_AD --datapath <data_path>
 ```
 **Tiny set formed by PTS (60 sheets) on MVTec AD:**
 ```bash
-python run.py --mode temp --ttype PTS --tsize 60 --dataset MVTec_AD
+python run.py --mode temp --ttype PTS --tsize 60 --dataset MVTec_AD --datapath <data_path>
 ```
 Since generating pixel-level OPTICS clusters is time-consuming, you can download the "*template*" folder from [Google Drive](https://drive.google.com/drive/folders/1c4XvmugX-ryP168bDMFcScdiYWgYktlu?usp=drive_link) / [Baidu Cloud](https://pan.baidu.com/s/1HH_3FQo1K72HbUvZpfylxw?pwd=eeg9) and copy it into our main folder as:
 ```
@@ -71,11 +71,11 @@ HETMM/
 ### 3) Anomaly Prediction
 **Original template set on MVTec AD:**
 ```bash
-python run.py --mode test --ttype ALL --dataset MVTec_AD
+python run.py --mode test --ttype ALL --dataset MVTec_AD --datapath <data_path>
 ```
 **Tiny set formed by PTS (60 sheets) on MVTec AD:**
 ```bash
-python run.py --mode test --ttype PTS --tsize 60 --dataset MVTec_AD
+python run.py --mode test --ttype PTS --tsize 60 --dataset MVTec_AD --datapath <data_path>
 ```
 Please see "*run.sh*" and "*run.py*" for more details.
 

configs/MVTec_AD.yaml

@@ -1,4 +1,12 @@
+dataset: MVTec_AD
 backbone: Wide_ResNet101_2
+temp_size: None
+temp_type: ALL
+test_batch_size: 32
+metric_functions : 
+- img_AUC
+- pix_AUC
+- PRO
 categories:
 - carpet
 - grid
@@ -15,7 +23,6 @@ categories:
 - toothbrush
 - transistor
 - zipper
-dataset: MVTec_AD
 mparams:
   bottle: L1:5x3_a=0.3_l=1|L2:5x3_a=0.3_l=1|L3:7x3_a=0.3_l=0.8;img_AUC:clean|pix_AUC:blur_s=4.1_k=39|PRO:clean
   cable: L1:13x7_a=0.7_l=1|L2:5x5_a=0.7_l=0.8|L3:7x1_a=0.3_l=0.8;img_AUC:blur_s=5.5_k=39|pix_AUC:blur_s=8.600000000000001_k=61|PRO:clean
@@ -31,12 +38,4 @@ mparams:
   toothbrush: L1:13x3_a=0.3_l=1|L2:9x5_a=0.3_l=0.8|L3:7x3_a=0.3_l=0.5;img_AUC:clean|pix_AUC:blur_s=3.0_k=11|PRO:clean
   transistor: L1:13x5_a=0.3_l=1|L2:5x3_a=0.5_l=0.5|L3:5x1_a=0.3_l=0.8;img_AUC:blur_s=5.1_k=31|pix_AUC:blur_s=15.0_k=61|PRO:blur_s=6.6_k=39
   wood: L1:7x3_a=0.7_l=1|L2:5x3_a=0.7_l=0.8|L3:7x1_a=0.7_l=0.3;img_AUC:clean|pix_AUC:clean|PRO:clean
-  zipper: L1:13x5_a=0.3_l=1|L2:7x5_a=0.7_l=1|L3:7x1_a=0.7_l=0.3;img_AUC:mp_k=11|pix_AUC:blur_s=10.4_k=61|PRO:blur_s=6.0_k=31
-root: /media/data4/chenzx/dataset/anomaly/mvtec_anomaly_detection
-temp_size: None
-temp_type: ALL
-test_batch_size: 32
-metric_functions : 
-- img_AUC
-- pix_AUC
-- PRO
+  zipper: L1:13x5_a=0.3_l=1|L2:7x5_a=0.7_l=1|L3:7x1_a=0.7_l=0.3;img_AUC:mp_k=11|pix_AUC:blur_s=10.4_k=61|PRO:blur_s=6.0_k=31

configs/base.yaml

@@ -1,6 +1,6 @@
 data : 
     name     : $dataset
-    root     : $root
+    datapath : $datapath
     out_size : $out_size
     normalize: 
         mean : [0.485, 0.456, 0.406]

run.py

@@ -16,6 +16,7 @@ def argParse():
     parser.add_argument('--method', default='ATMM')
     parser.add_argument('--ttype', choices=['ALL', 'PTS'], default='ALL')
     parser.add_argument('--tsize', type=int, default=0)
+    parser.add_argument('--datapath', help='your own data path')
     parser.add_argument('--dataset', type=str, default='MVTec_AD')
     parser.add_argument('--categories', type=str, nargs='+', default=None)
     parser.add_argument('--half', action='store_true')
@@ -55,7 +56,6 @@ def temp(cfg):
     def get_ALL(cfg, tpath):
         try:
             tdict = cfg.model.load_template(os.path.join(tpath, f'{cfg.model.backbone.lower()}_ALL.pkl'))
-            print ('Found the original template!')
 
         except:
             tdict = tl.gen_by_ALL(cfg.model, cfg.temploader, tpath, cfg.model.backbone.lower(), cfg.half, save=True)
@@ -68,7 +68,6 @@ def get_ALL(cfg, tpath):
     else:
         try:
             tdict = cfg.model.load_template(os.path.join(tpath, tname))
-            print (f'Found the {cfg.ttype}x{cfg.tsize}!')
 
         except:
             tdict = getattr(tl, f'gen_by_{cfg.ttype}')(get_ALL(cfg, tpath), cfg.tsize, tpath, cfg.model.backbone.lower(), num_workers=cfg.num_workers, save=True)
@@ -78,12 +77,9 @@ def get_ALL(cfg, tpath):
 if __name__ == '__main__':
     args = argParse()
     cfg = Cfg(args)
-    categories = tqdm(cfg.categories) if args.mode == 'test' else cfg.categories
+    categories = tqdm(cfg.categories)
     for category in categories:
-        if args.mode == 'test':
-            categories.set_description(category)
-        else:
-            print (category)
+        categories.set_description(category)
         cfg.update(category)
         globals()[args.mode](cfg)
 

run.sh

@@ -4,12 +4,12 @@ ts=60
 
 # Generating template
 # Original template
-python run.py --mode temp --ttype ALL --dataset $dataset
+python run.py --mode temp --ttype ALL --dataset $dataset --datapath <data_path>
 # PTS template
-python run.py --mode temp --ttype $tt --tsize $ts --dataset $dataset
+python run.py --mode temp --ttype $tt --tsize $ts --dataset $dataset --datapath <data_path>
 
 # Anomaly detection and localization
 # Original template
-python run.py --mode test --ttype ALL --dataset $dataset
+python run.py --mode test --ttype ALL --dataset $dataset --datapath <data_path>
 # PTS template
-python run.py --mode test --ttype $tt --tsize $ts --dataset $dataset
+python run.py --mode test --ttype $tt --tsize $ts --dataset $dataset --datapath <data_path>

src/config.py

@@ -38,10 +38,9 @@ def update(self, category):
             os.makedirs(os.path.join(getattr(self, key), self.dataset, self.category), exist_ok=True)
 
         for func in ['data', 'template']:
-            try:
-                getattr(self, f'load_{func}')()
-            except:
-                pass
+            getattr(self, f'load_{func}')()
+            # except:
+            #     pass
 
         self.model.update(category)
 

src/dataset.py

@@ -58,9 +58,9 @@ def __init__(self, **params):
         super(MVTec_AD, self).__init__(params)
 
     def load_data(self):
-        self.img_path = os.path.join(self.root, self.category, self.mode if self.mode == 'test' else 'train')
+        self.img_path = os.path.join(self.datapath, self.category, self.mode if self.mode == 'test' else 'train')
         if self.mode == 'test':
-            self.gt_path = os.path.join(self.root, self.category, 'ground_truth')
+            self.gt_path = os.path.join(self.datapath, self.category, 'ground_truth')
         self.img_paths, self.gt_paths, self.labels, self.types = [], [], [], []
         for defect_type in filter(lambda x : os.path.isdir(os.path.join(self.img_path, x)), os.listdir(self.img_path)):
             img_paths = [os.path.join(self.img_path, defect_type, x) for x in sorted(filter(lambda x : x.endswith('.png'), \

src/template.py

@@ -13,7 +13,7 @@
 
 def gen_by_ALL(model, temploader, tpath, backbone, half=False, save=True):
     Out_dict = {}
-    print (f'Generating the original template')
+    # print (f'Generating the original template')
     with torch.no_grad():
         for batch in tqdm(temploader):
             x = batch[0].cuda().half() if half else batch[0].cuda()
@@ -29,7 +29,7 @@ def gen_by_ALL(model, temploader, tpath, backbone, half=False, save=True):
 def gen_by_PC_OPTICS(tdict, tsize, tpath, backbone, num_workers, save=True, **kwargs):
     clu_path = os.path.join(tpath, f'{backbone}_PC_OPTICSx{tsize}.pkl')
     if os.path.exists(clu_path):
-        print ('Found pretrained OPTICS clusters!')
+        # print ('Found pretrained OPTICS clusters!')
         clu_dict = torch.load(clu_path, weights_only=True, map_location='cpu')
 
     else:
@@ -45,7 +45,7 @@ def gen_by_PTS(tdict, tsize, tpath, backbone, num_workers, save=True, **kwargs):
     clu_dict = gen_by_PC_OPTICS(tdict, tsize, tpath, backbone, num_workers, save=True, **kwargs)
 
     Out_dict = {}
-    print ('Generating PTS...')
+    # print ('Generating PTS...')
     for key, Ts in tqdm(tdict.items()):
         Out_dict[key] = PTS(Ts, clu_dict[key].to(Ts.device).long(), int(tsize))
 
@@ -62,7 +62,7 @@ def get_optics_clusters_unit(pixel, min_samples, metric):
 
     clu_dict = {}
 
-    print ('Generating OPTICS clusters (very slow)...')
+    # print ('Generating OPTICS clusters (very slow)...')
     for key, Ts in tqdm(tdict.items()):
         N, C, H, W = Ts.shape
         pixels = F.unfold(Ts, 1).cpu().numpy()