Update script.py

script.py

@@ -5,6 +5,7 @@ from PIL import Image
 from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
 from tqdm import tqdm
 
+
 def run_inference(image_path, model, save_path, prompt, box_threshold, text_threshold, device):
     
     try:
@@ -26,13 +27,17 @@ def run_inference(image_path, model, save_path, prompt, box_threshold, text_thre
         
         try:
             full_img_path = os.path.join(image_path, image_name)
-            img = Image.open(full_img_path).convert("RGB") 
-        except Exception as e:
+            img = Image.open(full_img_path).convert("RGB")
+        except Exception:
             bboxes.append([])
             category_ids.append([])
             continue
         
-        inputs = processor(images=img, text=prompt, return_tensors="pt").to(device)
+        inputs = processor(
+            images=img,
+            text=prompt,
+            return_tensors="pt"
+        ).to(device)
         
         with torch.no_grad():
             outputs = model(**inputs)
@@ -45,7 +50,7 @@ def run_inference(image_path, model, save_path, prompt, box_threshold, text_thre
             target_sizes=[img.size[::-1]]
         )
         
-        # Safe Mode: ID=0
+        # Safe Mode: Single category (ID = 0)
         for result in results:
             boxes = result["boxes"]
             for box in boxes:
@@ -53,7 +58,7 @@ def run_inference(image_path, model, save_path, prompt, box_threshold, text_thre
                 width = xmax - xmin
                 height = ymax - ymin
                 bbox.append([xmin, ymin, width, height])
-                category_id.append(0) 
+                category_id.append(0)
         
         bboxes.append(bbox)
         category_ids.append(category_id)
@@ -73,11 +78,13 @@ def run_inference(image_path, model, save_path, prompt, box_threshold, text_thre
 
 
 if __name__ == "__main__":
+    # Offline HuggingFace settings
     os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
     os.environ["HF_HUB_OFFLINE"] = "1"
     os.environ["HF_DATASETS_OFFLINE"] = "1"
     
     current_directory = os.path.dirname(os.path.abspath(__file__))
+    
     TEST_IMAGE_PATH = "/tmp/data/test_images"
     SUBMISSION_SAVE_PATH = os.path.join(current_directory, "submission.csv")
     
@@ -89,13 +96,38 @@ if __name__ == "__main__":
     processor = AutoProcessor.from_pretrained(processor_path)
     model = AutoModelForZeroShotObjectDetection.from_pretrained(model_path)
     model.to(device)
+    model.eval()
     
-    # --- OPTIMIZED SETTINGS ---
-    # 1. Prompt: Reverting to the one that got 0.047 ("Large" helps!)
-    PROMPT = "Monopolar Curved Scissors . Prograsp Forceps . Large Needle Driver ."
+    # =========================
+    # 🔥 PROMPT ENGINEERING
+    # =========================
+    PROMPT = (
+        "Monopolar Curved Scissors. "
+        "curved surgical scissors. "
+        "surgical scissors cutting tissue. "
+        "Prograsp Forceps. "
+        "surgical forceps grasping tissue. "
+        "grasping forceps. "
+        "Large Needle Driver. "
+        "needle holder. "
+        "surgical needle driver. "
+        "laparoscopic surgical instrument. "
+        "robotic surgical instrument. "
+        "metal surgical tool inside the body."
+    )
     
-    # 2. Threshold: Increased to 0.35 (Stricter than 0.25 to improve Precision)
-    BOX_THRESHOLD = 0.35
-    TEXT_THRESHOLD = 0.25
+    # =========================
+    # 🎯 THRESHOLDS (Recall-Oriented)
+    # =========================
+    BOX_THRESHOLD = 0.25
+    TEXT_THRESHOLD = 0.20
     
-    run_inference(TEST_IMAGE_PATH, model, SUBMISSION_SAVE_PATH, PROMPT, BOX_THRESHOLD, TEXT_THRESHOLD, device)
+    run_inference(
+        TEST_IMAGE_PATH,
+        model,
+        SUBMISSION_SAVE_PATH,
+        PROMPT,
+        BOX_THRESHOLD,
+        TEXT_THRESHOLD,
+        device
+    )