Spaces:
Running
on
Zero
Running
on
Zero
Martin Tomov
commited on
cv2.rectangle attempt
Browse files
app.py
CHANGED
@@ -25,7 +25,6 @@ class BoundingBox:
|
|
25 |
@property
|
26 |
def xyxy(self) -> List[float]:
|
27 |
return [self.xmin, self.ymin, self.xmax, self.ymax]
|
28 |
-
|
29 |
@dataclass
|
30 |
class DetectionResult:
|
31 |
score: float
|
@@ -49,24 +48,31 @@ 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 |
-
|
59 |
-
|
60 |
-
cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color, 2)
|
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 |
mask_uint8 = (mask * 255).astype(np.uint8)
|
66 |
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
67 |
-
cv2.drawContours(
|
|
|
|
|
|
|
|
|
68 |
|
69 |
-
|
|
|
|
|
|
|
|
|
|
|
70 |
|
71 |
def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult]) -> np.ndarray:
|
72 |
annotated_image = annotate(image, detections)
|
@@ -106,7 +112,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]
|
|
|
25 |
@property
|
26 |
def xyxy(self) -> List[float]:
|
27 |
return [self.xmin, self.ymin, self.xmax, self.ymax]
|
|
|
28 |
@dataclass
|
29 |
class DetectionResult:
|
30 |
score: float
|
|
|
48 |
def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[DetectionResult]) -> np.ndarray:
|
49 |
image_cv2 = np.array(image) if isinstance(image, Image.Image) else image
|
50 |
image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)
|
51 |
+
|
52 |
+
# Make the entire background yellow
|
53 |
+
yellow_background = np.full(image_cv2.shape, (0, 255, 255), dtype=np.uint8)
|
54 |
+
|
55 |
for detection in detection_results:
|
56 |
label = detection.label
|
57 |
score = detection.score
|
58 |
box = detection.box
|
59 |
mask = detection.mask
|
60 |
+
|
|
|
|
|
|
|
|
|
|
|
61 |
if mask is not None:
|
62 |
mask_uint8 = (mask * 255).astype(np.uint8)
|
63 |
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
64 |
+
cv2.drawContours(yellow_background, contours, -1, (0, 0, 0), cv2.FILLED)
|
65 |
+
|
66 |
+
# Drawing bounding box (border and inside should be yellow)
|
67 |
+
cv2.rectangle(yellow_background, (box.xmin, box.ymin), (box.xmax, box.ymax), (0, 255, 255), cv2.FILLED) # Fill inside with yellow
|
68 |
+
cv2.rectangle(yellow_background, (box.xmin, box.ymin), (box.xmax, box.ymax), (0, 255, 255), 2) # Draw border yellow
|
69 |
|
70 |
+
if mask is not None:
|
71 |
+
# Overlay the insect on top of the yellow background
|
72 |
+
insect_region = image_cv2[box.ymin:box.ymax, box.xmin:box.xmax]
|
73 |
+
yellow_background[box.ymin:box.ymax, box.xmin:box.xmax] = cv2.bitwise_and(yellow_background[box.ymin:box.ymax, box.xmin:box.xmax], insect_region, mask=mask[box.ymin:box.ymax, box.xmin:box.xmax])
|
74 |
+
|
75 |
+
return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
|
76 |
|
77 |
def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult]) -> np.ndarray:
|
78 |
annotated_image = annotate(image, detections)
|
|
|
112 |
return list(masks)
|
113 |
|
114 |
@spaces.GPU
|
115 |
+
def detect(image: Image.Image, labels: List[str], threshold: float = 0.3, detector_id: Optional[str] = None) -> List<Dict[str, Any]]:
|
116 |
detector_id = detector_id if detector_id else "IDEA-Research/grounding-dino-base"
|
117 |
object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device="cuda")
|
118 |
labels = [label if label.endswith(".") else label+"." for label in labels]
|