BUG FIX

CASIA-IVA-Lab · Sep 1, 2023 · ac4164b · ac4164b
1 parent 59fab40
commit ac4164b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 12 deletions.
diff --git a/code/datasets/mvtec.py b/code/datasets/mvtec.py
@@ -191,7 +191,7 @@ def __getitem__(self, index):
 
             conversation_abnormal = conversation_normal
 
-        return origin, conversation_normal, x, conversation_abnormal, class_name, mask
+        return origin, conversation_normal, x, conversation_abnormal, class_name, mask, img_path
 
 
 
@@ -201,20 +201,24 @@ def collate(self, instances):
         texts = []
         class_names = []
         masks = []
+        img_paths = []
         for instance in instances:
             images.append(instance[0])
             texts.append(instance[1])
             class_names.append(instance[4])
             masks.append(torch.zeros_like(instance[5]))
+            img_paths.append(instance[6])
 
             images.append(instance[2])
             texts.append(instance[3])
             class_names.append(instance[4])
             masks.append(instance[5])
+            img_paths.append(instance[6])
 
         return dict(
             images=images,
             texts=texts,
             class_names=class_names,
-            masks=masks
+            masks=masks,
+            img_paths=img_paths
         )
diff --git a/code/datasets/visa.py b/code/datasets/visa.py
@@ -172,7 +172,7 @@ def __getitem__(self, index):
             conversation_abnormal = conversation_normal
 
 
-        return origin, conversation_normal, x, conversation_abnormal, class_name, mask
+        return origin, conversation_normal, x, conversation_abnormal, class_name, mask, img_path
 
 
 
@@ -182,21 +182,25 @@ def collate(self, instances):
         texts = []
         class_names = []
         masks = []
+        img_paths = []
         for instance in instances:
             images.append(instance[0])
             texts.append(instance[1])
             class_names.append(instance[4])
             masks.append(torch.zeros_like(instance[5]))
+            img_paths.append(instance[6])
 
             images.append(instance[2])
             texts.append(instance[3])
             class_names.append(instance[4])
             masks.append(instance[5])
+            img_paths.append(instance[6])
 
 
         return dict(
             images=images,
             texts=texts,
             class_names=class_names,
-            masks=masks
+            masks=masks,
+            img_paths=img_paths
         )
diff --git a/code/model/openllama.py b/code/model/openllama.py
@@ -443,22 +443,18 @@ def forward(self, inputs):
             anomaly_maps = []
             for layer in range(len(patch_tokens)):
                 patch_tokens[layer] = patch_tokens[layer] / patch_tokens[layer].norm(dim=-1, keepdim=True)
-                # print(patch_tokens[layer].shape)
-                # anomaly_map = torch.bmm(patch_tokens[layer], feats_text_tensor.transpose(-2,-1))
                 anomaly_map = (100.0 * patch_tokens[layer] @ feats_text_tensor.transpose(-2,-1))
                 B, L, C = anomaly_map.shape
                 H = int(np.sqrt(L))
                 anomaly_map = F.interpolate(anomaly_map.permute(0, 2, 1).view(B, 2, H, H),
                                             size=224, mode='bilinear', align_corners=True)
-                # anomaly_map_no_softmax = anomaly_map
+
                 anomaly_map = torch.softmax(anomaly_map, dim=1)
-                anomaly_maps.append(anomaly_map)
-                # anomaly_maps_ns.append(anomaly_map_no_softmax)           
+                anomaly_maps.append(anomaly_map)           
 
             gt = inputs['masks']
             gt = torch.stack(gt, dim=0).to(self.device)
             gt = gt.squeeze()
-            # print(gt.max(), gt.min())
             gt[gt > 0.3], gt[gt <= 0.3] = 1, 0
 
 
@@ -476,8 +472,12 @@ def forward(self, inputs):
 
                 normal_paths = []
                 for path in inputs['img_paths']:
-                    normal_path = path.replace('test', 'train')
-                    normal_path = find_first_file_in_directory("/".join(normal_path.split('/')[:-2])+'/good')
+                    if 'visa' in image_paths.lower():
+                        normal_path = path.replace('Anomaly', 'Normal')
+                        normal_path = find_first_file_in_directory("/".join(normal_path.split('/')[:-1]))
+                    else:
+                        normal_path = path.replace('test', 'train')
+                        normal_path = find_first_file_in_directory("/".join(normal_path.split('/')[:-2])+'/good')
                     normal_paths.append(normal_path)
 
                 print(normal_paths)