optimise

app.py

@@ -9,10 +9,8 @@ import torch
 import requests
 import numpy as np
 from PIL import Image
-import matplotlib.pyplot as plt
-from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
 import gradio as gr
-import spaces
+from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
 import json
 
 @dataclass
@@ -54,10 +52,9 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
         label = detection.label
         score = detection.score
         box = detection.box
-        mask = detection.mask
 
         if include_bboxes:
-            color = np.random.randint(0, 256, size=3).tolist()
+            color = [int(c) for c in np.random.randint(0, 256, size=3)]
             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)
@@ -65,8 +62,7 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
     return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
 
 def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:
-    annotated_image = annotate(image, detections, include_bboxes)
-    return annotated_image
+    return annotate(image, detections, include_bboxes)
 
 def load_image(image: Union[str, Image.Image]) -> Image.Image:
     if isinstance(image, str) and image.startswith("http"):
@@ -77,19 +73,14 @@ def load_image(image: Union[str, Image.Image]) -> Image.Image:
         image = image.convert("RGB")
     return image
 
-def get_boxes(detection_results: List[DetectionResult]) -> List[List[List[float]]]:
-    boxes = []
-    for result in detection_results:
-        xyxy = result.box.xyxy
-        boxes.append(xyxy)
-    return [boxes]
+def get_boxes(detection_results: List[DetectionResult]) -> List[List[float]]:
+    return [result.box.xyxy for result in detection_results]
 
 def mask_to_polygon(mask: np.ndarray) -> np.ndarray:
     contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if len(contours) == 0:
         return np.array([])
-    largest_contour = max(contours, key=cv2.contourArea)
-    return largest_contour
+    return max(contours, key=cv2.contourArea)
 
 def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]:
     masks = masks.cpu().float().permute(0, 2, 3, 1).mean(axis=-1).numpy().astype(np.uint8)
@@ -101,21 +92,19 @@ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> L
             masks[idx] = cv2.fillPoly(np.zeros(shape, dtype=np.uint8), [polygon], 1)
     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]]:
+def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[DetectionResult]:
     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="cuda")
-    labels = [label if label.endswith(".") else label+"." for label in labels]
+    object_detector = pipeline(model=detector_id, task="zero-shot-object-detection")
+    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]:
     segmenter_id = segmenter_id if segmenter_id else "martintmv/InsectSAM"
-    segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to("cuda")
+    segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id)
     processor = AutoProcessor.from_pretrained(segmenter_id)
     boxes = get_boxes(detection_results)
-    inputs = processor(images=image, input_boxes=boxes, return_tensors="pt").to("cuda")
+    inputs = processor(images=image, input_boxes=boxes, return_tensors="pt")
     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)
@@ -152,9 +141,7 @@ def create_yellow_background_with_insects(image: np.ndarray, detections: List[De
     for detection in detections:
         if detection.mask is not None:
             extract_and_paste_insect(image, detection, yellow_background)
-    # Convert back to RGB to match Gradio's expected input format
-    yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
-    return yellow_background
+    return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
 
 def run_length_encoding(mask):
     pixels = mask.flatten()