Spaces:
Running
on
Zero
Running
on
Zero
Martin Tomov
commited on
millionth experiment with bbox
Browse files
app.py
CHANGED
@@ -49,24 +49,19 @@ class DetectionResult:
|
|
49 |
def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult]) -> np.ndarray:
|
50 |
image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
|
51 |
image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
|
52 |
-
|
|
|
|
|
|
|
53 |
for detection in detection_results:
|
54 |
-
label = detection.label
|
55 |
-
score = detection.score
|
56 |
-
box = detection.box
|
57 |
mask = detection.mask
|
58 |
-
color = (0, 255, 0) # Green fill color for debugging
|
59 |
|
60 |
-
cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), (0, 0, 255), -1)
|
61 |
-
cv2.putText(image_cv2, f'{label}: {score:.2f}', (box.xmin, box.ymin - 10),
|
62 |
-
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
63 |
-
|
64 |
if mask is not None:
|
65 |
-
|
66 |
-
|
67 |
-
|
68 |
|
69 |
-
return cv2.cvtColor(
|
70 |
|
71 |
def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult]) -> np.ndarray:
|
72 |
annotated_image = annotate(image, detections)
|
@@ -106,7 +101,7 @@ def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> L
|
|
106 |
return list(masks)
|
107 |
|
108 |
@spaces.GPU
|
109 |
-
def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[
|
110 |
detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
|
111 |
object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device="cuda")
|
112 |
labels = [label if label.endswith(".") else label+"." for label in labels]
|
@@ -159,7 +154,7 @@ def create_yellow_background_with_insects(image: np.ndarray, detections: List[De
|
|
159 |
yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
|
160 |
return yellow_background
|
161 |
|
162 |
-
def run_length_encoding(mask):
|
163 |
pixels = mask.flatten()
|
164 |
rle = []
|
165 |
last_val = 0
|
@@ -176,7 +171,7 @@ def run_length_encoding(mask):
|
|
176 |
rle.append(count)
|
177 |
return rle
|
178 |
|
179 |
-
def detections_to_json(detections):
|
180 |
detections_list = []
|
181 |
for detection in detections:
|
182 |
detection_dict = {
|
|
|
49 |
def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult]) -> np.ndarray:
|
50 |
image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
|
51 |
image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
|
52 |
+
|
53 |
+
# Make the entire background yellow
|
54 |
+
yellow_background = np.full(image_cv2.shape, (0, 255, 255), dtype=np.uint8)
|
55 |
+
|
56 |
for detection in detection_results:
|
|
|
|
|
|
|
57 |
mask = detection.mask
|
|
|
58 |
|
|
|
|
|
|
|
|
|
59 |
if mask is not None:
|
60 |
+
mask_expanded = np.stack([mask]*3, axis=-1) # Expand mask dimensions for color channels
|
61 |
+
insect_region = np.where(mask_expanded, image_cv2, yellow_background)
|
62 |
+
yellow_background = np.where(mask_expanded, insect_region, yellow_background)
|
63 |
|
64 |
+
return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
|
65 |
|
66 |
def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult]) -> np.ndarray:
|
67 |
annotated_image = annotate(image, detections)
|
|
|
101 |
return list(masks)
|
102 |
|
103 |
@spaces.GPU
|
104 |
+
def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List[DetectionResult]:
|
105 |
detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
|
106 |
object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device="cuda")
|
107 |
labels = [label if label.endswith(".") else label+"." for label in labels]
|
|
|
154 |
yellow_background = cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
|
155 |
return yellow_background
|
156 |
|
157 |
+
def run_length_encoding(mask: np.ndarray) -> List[int]:
|
158 |
pixels = mask.flatten()
|
159 |
rle = []
|
160 |
last_val = 0
|
|
|
171 |
rle.append(count)
|
172 |
return rle
|
173 |
|
174 |
+
def detections_to_json(detections: List[DetectionResult]) -> List[Dict[str, Any]]:
|
175 |
detections_list = []
|
176 |
for detection in detections:
|
177 |
detection_dict = {
|