Update sam_utils.py

sam_utils.py

@@ -1,5 +1,5 @@
 import os
-import spaces
+
 import random
 from dataclasses import dataclass
 from typing import Any, List, Dict, Optional, Union, Tuple
@@ -13,6 +13,7 @@ from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
 import gradio as gr
 import json
 
+
 @dataclass
 class BoundingBox:
     xmin: int
@@ -23,7 +24,6 @@ class BoundingBox:
     @property
     def xyxy(self) -> List[float]:
         return [self.xmin, self.ymin, self.xmax, self.ymax]
-
 @dataclass
 class DetectionResult:
     score: float
@@ -63,10 +63,12 @@ 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")
@@ -76,6 +78,7 @@ 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:
@@ -83,6 +86,7 @@ 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)
@@ -91,6 +95,7 @@ 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)
@@ -103,7 +108,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(
@@ -113,7 +118,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(
@@ -130,16 +135,19 @@ 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)
@@ -154,6 +162,7 @@ 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"]
 
@@ -170,6 +179,7 @@ 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 = []
@@ -187,6 +197,7 @@ def run_length_encoding(mask):
         rle.append(count)
     return rle
 
+
 def detections_to_json(detections):
     detections_list = []
     for detection in detections:
@@ -203,6 +214,7 @@ 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: