fix categories

app.py

@@ -31,7 +31,6 @@ from detectron2.checkpoint import DetectionCheckpointer
 from demo.demo import VisualizationDemo, get_parser
 import numpy as np
 from detectron2.data.detection_utils import read_image
-from detrex.data.datasets.register_ade20k_full import _get_ade20k_full_meta
 import matplotlib.pyplot as plt
 import cv2
 
@@ -40,6 +39,7 @@ warnings.filterwarnings("ignore")
 
 config_file = os.getcwd() + '/projects/dino/configs/odor3_fn_l_lrf_384_fl4_5scale_50ep.py'
 ckpt_pth = os.getcwd() + '/utils/focaldino_ep18.pth'
+categories = read_json_categories(os.getcwd() + '/annotations/instances_train2017.json')
 
 try:
     cfg = LazyConfig.load(config_file)
@@ -60,30 +60,25 @@ demo = VisualizationDemo(
     img_format='RGB',
     metadata_dataset='odor_test')
 
+def read_json_categories(jsonFile):
+    categories_dict = {}
+    with open(jsonFile, 'r') as file:
+        data = json.load(file)
+        if 'categories' in data:
+            categories_dict = data['categories']
+    return categories_dict
+
 def treat_grayscale(img):
     if len(img.shape) == 2:
         return np.stack((img,)*3, axis=-1)
     else:
         return img
 
-def get_name_by_id(contiguous_id, ade20k_meta):
-    stuff_dataset_id_to_contiguous_id = ade20k_meta["stuff_dataset_id_to_contiguous_id"]
-    stuff_classes = ade20k_meta["stuff_classes"]
-    dataset_id = -1
-    print("stuff:")
-    print(stuff_classes)
-    print("con id:")
-    print(contiguous_id)
-    if contiguous_id in stuff_dataset_id_to_contiguous_id:
-        dataset_id = stuff_dataset_id_to_contiguous_id[contiguous_id]
-        print("data id:")
-        print(dataset_id)
-        
-    if 0 <= dataset_id < len(stuff_classes):
-        class_name = stuff_classes[dataset_id]
-        return class_name
-        
-    return "Unknown"
+def get_name_by_id(categories, id):
+    for cg in categories:
+        if cg['id'] == id:
+            return cg['name']
+    return 'Unknown'
 
 def predict(link, text, threshold):
     if(link):
@@ -102,15 +97,15 @@ def predict(link, text, threshold):
 
     output_text = ""
     for i in range(len(pred_boxes)):
-        continuous_id = pred_classes[i].item()
-        class_name = get_name_by_id(continuous_id, _get_ade20k_full_meta())
+        id = pred_classes[i].item()
+        class_name = get_name_by_id(categories, id)
         score = scores[i].item()
         output_text += f"{class_name}: {score:.4f}\n"
 
     output_json = []
     for i in range(len(pred_boxes)):
-        continuous_id = pred_classes[i].item()
-        class_name = get_name_by_id(continuous_id, _get_ade20k_full_meta())
+        id = pred_classes[i].item()
+        class_name = get_name_by_id(categories, id)
         score = scores[i].item()
         box_coords = pred_boxes[i].tensor.tolist()
         output_json.append({
@@ -119,7 +114,6 @@ def predict(link, text, threshold):
             "box_coordinates": box_coords
         })
     output_json = json.dumps(output_json, indent=4)
-    output_json = ""
     
     return visualized_output.get_image(), output_text, output_json