spaces.GPU

app.py

@@ -92,19 +92,21 @@ 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[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")
+    object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device=0)
     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)
+    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")
+    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)