update app.py to match ZeroGPU Spaces

app.py

@@ -1,7 +1,6 @@
 import random
 from dataclasses import dataclass
 from typing import Any, List, Dict, Optional, Union, Tuple
-
 import cv2
 import torch
 import requests
@@ -10,6 +9,7 @@ from PIL import Image
 import matplotlib.pyplot as plt
 from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
 import gradio as gr
+import spaces
 
 @dataclass
 class BoundingBox:
@@ -31,12 +31,16 @@ class DetectionResult:
 
     @classmethod
     def from_dict(cls, detection_dict: Dict) -> 'DetectionResult':
-        return cls(score=detection_dict['score'],
-                   label=detection_dict['label'],
-                   box=BoundingBox(xmin=detection_dict['box']['xmin'],
-                                   ymin=detection_dict['box']['ymin'],
-                                   xmax=detection_dict['box']['xmax'],
-                                   ymax=detection_dict['box']['ymax']))
+        return cls(
+            score=detection_dict['score'],
+            label=detection_dict['label'],
+            box=BoundingBox(
+                xmin=detection_dict['box']['xmin'],
+                ymin=detection_dict['box']['ymin'],
+                xmax=detection_dict['box']['xmax'],
+                ymax=detection_dict['box']['ymax']
+            )
+        )
 
 def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult]) -> np.ndarray:
     image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
@@ -47,16 +51,16 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
         score = detection.score
         box = detection.box
         mask = detection.mask
+        color = np.random.randint(0, 256, size=3).tolist()
 
-        color = np.random.randint(0, 256, size=3)
-
-        cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color.tolist(), 2)
-        cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color.tolist(), 2)
+        cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
+        cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
 
         if mask is not None:
             mask_uint8 = (mask * 255).astype(np.uint8)
             contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            cv2.drawContours(image_cv2, contours, -1, color.tolist(), 2)
+            cv2.drawContours(image_cv2, contours, -1, color, 2)
 
     return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
 
@@ -90,29 +94,26 @@ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> L
             masks[idx] = polygon_to_mask(polygon, shape)
     return list(masks)
 
+@spaces.GPU
 def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[Dict[str, Any]]:
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
-    object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device=device)
+    object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device="cuda")
     labels = [label if label.endswith(".") else label+"." for label in labels]
     results = object_detector(image, candidate_labels=labels, threshold=threshold)
     return [DetectionResult.from_dict(result) for result in results]
 
+@spaces.GPU
 def segment(image: Image.Image, detection_results: List[DetectionResult], polygon_refinement: bool = False, segmenter_id: Optional[str] = None) -> List[DetectionResult]:
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     segmenter_id = segmenter_id if segmenter_id else "martintmv/InsectSAM"
-    segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(device)
+    segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to("cuda")
     processor = AutoProcessor.from_pretrained(segmenter_id)
-
     boxes = get_boxes(detection_results)
-    inputs = processor(images=image, input_boxes=boxes, return_tensors="pt").to(device)
+    inputs = processor(images=image, input_boxes=boxes, return_tensors="pt").to("cuda")
     outputs = segmentator(**inputs)
     masks = processor.post_process_masks(masks=outputs.pred_masks, original_sizes=inputs.original_sizes, reshaped_input_sizes=inputs.reshaped_input_sizes)[0]
     masks = refine_masks(masks, polygon_refinement)
-
     for detection_result, mask in zip(detection_results, masks):
         detection_result.mask = mask
-
     return detection_results
 
 def grounded_segmentation(image: Union[Image.Image, str], labels: List[str], threshold: float = 0.3, polygon_refinement: bool = False, detector_id: Optional[str] = None, segmenter_id: Optional[str] = None) -> Tuple[np.ndarray, List[DetectionResult]]:
@@ -165,8 +166,10 @@ def draw_classification_boxes(image_with_insects: np.ndarray, detections: List[D
         color = np.random.randint(0, 256, size=3).tolist()
         cv2.rectangle(image_with_insects, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
         (text_width, text_height), baseline = cv2.getTextSize(f"{label}: {score:.2f}", cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
-        cv2.rectangle(image_with_insects, (box.xmin, box.ymin - text_height - baseline), (box.xmin + text_width, box.ymin), color, thickness=cv2.FILLED)
-        cv2.putText(image_with_insects, f"{label}: {score:.2f}", (box.xmin, box.ymin - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)
+        cv2.rectangle(image_with_insects, (box.xmin, box.ymin - text_height - baseline), (box.xmin + text_width, box.ymin),
+                      color, thickness=cv2.FILLED)
+        cv2.putText(image_with_insects, f"{label}: {score:.2f}", (box.xmin, box.ymin - baseline),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)
     return image_with_insects
 
 def process_image(image):
@@ -177,7 +180,6 @@ def process_image(image):
     yellow_background_with_masks = put_masks_on_yellow_background(original_image.shape[:2], insect_masks)
     yellow_background_with_insects = create_yellow_background_with_insects(original_image, detections)
     yellow_background_with_boxes = draw_classification_boxes(yellow_background_with_insects, detections)
-
     return masked_image, yellow_background_with_masks, yellow_background_with_boxes
 
 gr.Interface(
@@ -185,4 +187,4 @@ gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs=[gr.Image(type="numpy"), gr.Image(type="numpy"), gr.Image(type="numpy")],
     title="Insect Detection and Masking"
-).launch()
+).launch()