Spaces:
Running
on
Zero
Running
on
Zero
Martin Tomov
commited on
Update app.py
Browse files
app.py
CHANGED
@@ -9,10 +9,11 @@ import torch
|
|
9 |
import requests
|
10 |
import numpy as np
|
11 |
from PIL import Image
|
12 |
-
import
|
13 |
from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
|
14 |
-
import
|
15 |
import spaces
|
|
|
16 |
|
17 |
@dataclass
|
18 |
class BoundingBox:
|
@@ -53,9 +54,10 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
|
|
53 |
label = detection.label
|
54 |
score = detection.score
|
55 |
box = detection.box
|
|
|
56 |
|
57 |
if include_bboxes:
|
58 |
-
color =
|
59 |
cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
|
60 |
cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
|
61 |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
@@ -63,7 +65,8 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
|
|
63 |
return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
|
64 |
|
65 |
def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:
|
66 |
-
|
|
|
67 |
|
68 |
def load_image(image: Union[str, Image.Image]) -> Image.Image:
|
69 |
if isinstance(image, str) and image.startswith("http"):
|
@@ -74,14 +77,19 @@ def load_image(image: Union[str, Image.Image]) -> Image.Image:
|
|
74 |
image = image.convert("RGB")
|
75 |
return image
|
76 |
|
77 |
-
def get_boxes(detection_results: List[DetectionResult]) -> List[List[float]]:
|
78 |
-
|
|
|
|
|
|
|
|
|
79 |
|
80 |
def mask_to_polygon(mask: np.ndarray) -> np.ndarray:
|
81 |
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
82 |
if len(contours) == 0:
|
83 |
return np.array([])
|
84 |
-
|
|
|
85 |
|
86 |
def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]:
|
87 |
masks = masks.cpu().float().permute(0, 2, 3, 1).mean(axis=-1).numpy().astype(np.uint8)
|
@@ -94,10 +102,10 @@ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> L
|
|
94 |
return list(masks)
|
95 |
|
96 |
@spaces.GPU
|
97 |
-
def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[
|
98 |
detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
|
99 |
-
object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device=
|
100 |
-
labels = [label if label.endswith(".") else label
|
101 |
results = object_detector(image, candidate_labels=labels, threshold=threshold)
|
102 |
return [DetectionResult.from_dict(result) for result in results]
|
103 |
|
@@ -144,7 +152,9 @@ def create_yellow_background_with_insects(image: np.ndarray, detections: List[De
|
|
144 |
for detection in detections:
|
145 |
if detection.mask is not None:
|
146 |
extract_and_paste_insect(image, detection, yellow_background)
|
147 |
-
|
|
|
|
|
148 |
|
149 |
def run_length_encoding(mask):
|
150 |
pixels = mask.flatten()
|
@@ -261,4 +271,4 @@ with gr.Blocks(css=css) as demo:
|
|
261 |
|
262 |
submit_button.click(update_outputs, [image_input, include_json, include_bboxes], [annotated_output, json_output, crops_output])
|
263 |
|
264 |
-
demo.launch()
|
|
|
9 |
import requests
|
10 |
import numpy as np
|
11 |
from PIL import Image
|
12 |
+
import matplotlib.pyplot as plt
|
13 |
from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
|
14 |
+
import gradio as gr
|
15 |
import spaces
|
16 |
+
import json
|
17 |
|
18 |
@dataclass
|
19 |
class BoundingBox:
|
|
|
54 |
label = detection.label
|
55 |
score = detection.score
|
56 |
box = detection.box
|
57 |
+
mask = detection.mask
|
58 |
|
59 |
if include_bboxes:
|
60 |
+
color = np.random.randint(0, 256, size=3).tolist()
|
61 |
cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
|
62 |
cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
|
63 |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
|
|
65 |
return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
|
66 |
|
67 |
def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult], include_bboxes: bool = True) -> np.ndarray:
|
68 |
+
annotated_image = annotate(image, detections, include_bboxes)
|
69 |
+
return annotated_image
|
70 |
|
71 |
def load_image(image: Union[str, Image.Image]) -> Image.Image:
|
72 |
if isinstance(image, str) and image.startswith("http"):
|
|
|
77 |
image = image.convert("RGB")
|
78 |
return image
|
79 |
|
80 |
+
def get_boxes(detection_results: List[DetectionResult]) -> List[List[List[float]]]:
|
81 |
+
boxes = []
|
82 |
+
for result in detection_results:
|
83 |
+
xyxy = result.box.xyxy
|
84 |
+
boxes.append(xyxy)
|
85 |
+
return [boxes]
|
86 |
|
87 |
def mask_to_polygon(mask: np.ndarray) -> np.ndarray:
|
88 |
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
89 |
if len(contours) == 0:
|
90 |
return np.array([])
|
91 |
+
largest_contour = max(contours, key=cv2.contourArea)
|
92 |
+
return largest_contour
|
93 |
|
94 |
def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]:
|
95 |
masks = masks.cpu().float().permute(0, 2, 3, 1).mean(axis=-1).numpy().astype(np.uint8)
|
|
|
102 |
return list(masks)
|
103 |
|
104 |
@spaces.GPU
|
105 |
+
def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[Dict[str, Any]]:
|
106 |
detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
|
107 |
+
object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device="cuda")
|
108 |
+
labels = [label if label.endswith(".") else label+"." for label in labels]
|
109 |
results = object_detector(image, candidate_labels=labels, threshold=threshold)
|
110 |
return [DetectionResult.from_dict(result) for result in results]
|
111 |
|
|
|
152 |
for detection in detections:
|
153 |
if detection.mask is not None:
|
154 |
extract_and_paste_insect(image, detection, yellow_background)
|
155 |
+
# Convert back to RGB to match Gradio's expected input format
|
156 |
+
yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
|
157 |
+
return yellow_background
|
158 |
|
159 |
def run_length_encoding(mask):
|
160 |
pixels = mask.flatten()
|
|
|
271 |
|
272 |
submit_button.click(update_outputs, [image_input, include_json, include_bboxes], [annotated_output, json_output, crops_output])
|
273 |
|
274 |
+
demo.launch()
|