@spaces.GPU sam_utils.py

sam_utils.py

@@ -1,5 +1,4 @@
 import os
-
 import random
 from dataclasses import dataclass
 from typing import Any, List, Dict, Optional, Union, Tuple
@@ -12,7 +11,7 @@ import matplotlib.pyplot as plt
 from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
 import gradio as gr
 import json
-
+import spaces
 
 @dataclass
 class BoundingBox:
@@ -24,6 +23,7 @@ class BoundingBox:
     @property
     def xyxy(self) -> List[float]:
         return [self.xmin, self.ymin, self.xmax, self.ymax]
+
 @dataclass
 class DetectionResult:
     score: float
@@ -63,12 +63,10 @@ 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
 
-
 def load_image(image: Union[str, Image.Image]) -> Image.Image:
     if isinstance(image, str) and image.startswith("http"):
         image = Image.open(requests.get(image, stream=True).raw).convert("RGB")
@@ -78,7 +76,6 @@ 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:
@@ -86,7 +83,6 @@ def get_boxes(detection_results: List[DetectionResult]) -> List[List[List[float]
         boxes.append(xyxy)
     return [boxes]
 
-
 def mask_to_polygon(mask: np.ndarray) -> np.ndarray:
     contours, _ = cv2.findContours(
         mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
@@ -95,7 +91,6 @@ def mask_to_polygon(mask: np.ndarray) -> np.ndarray:
     largest_contour = max(contours, key=cv2.contourArea)
     return largest_contour
 
-
 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)
@@ -108,7 +103,7 @@ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> L
                 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]]:
     detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
     object_detector = pipeline(
@@ -118,7 +113,7 @@ def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detect
         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(
@@ -135,19 +130,16 @@ def segment(image: Image.Image, detection_results: List[DetectionResult], polygo
         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]]:
     image = load_image(image)
     detections = detect(image, labels, threshold, detector_id)
     detections = segment(image, detections, polygon_refinement, segmenter_id)
     return np.array(image), detections
 
-
 def mask_to_min_max(mask: np.ndarray) -> Tuple[int, int, int, int]:
     y, x = np.where(mask)
     return x.min(), y.min(), x.max(), y.max()
 
-
 def extract_and_paste_insect(original_image: np.ndarray, detection: DetectionResult, background: np.ndarray) -> None:
     mask = detection.mask
     xmin, ymin, xmax, ymax = mask_to_min_max(mask)
@@ -162,7 +154,6 @@ def extract_and_paste_insect(original_image: np.ndarray, detection: DetectionRes
     insect_area = background[y_offset:y_end, x_offset:x_end]
     insect_area[mask_crop == 1] = insect[mask_crop == 1]
 
-
 def create_yellow_background_with_insects(image: np.ndarray) -> np.ndarray:
     labels = ["insect"]
 
@@ -179,14 +170,13 @@ def create_yellow_background_with_insects(image: np.ndarray) -> np.ndarray:
     yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
     return yellow_background
 
-
 def run_length_encoding(mask):
     pixels = mask.flatten()
     rle = []
     last_val = 0
     count = 0
     for pixel in pixels:
-        if pixel == last_val:
+        if pixel was the last val:
             count += 1
         else:
             if count > 0:
@@ -197,7 +187,6 @@ def run_length_encoding(mask):
         rle.append(count)
     return rle
 
-
 def detections_to_json(detections):
     detections_list = []
     for detection in detections:
@@ -214,7 +203,6 @@ def detections_to_json(detections):
         detections_list.append(detection_dict)
     return detections_list
 
-
 def crop_bounding_boxes_with_yellow_background(image: np.ndarray, yellow_background: np.ndarray, detections: List[DetectionResult]) -> List[np.ndarray]:
     crops = []
     for detection in detections: