Martin Tomov commited on
Commit
e1ab1fa
·
verified ·
1 Parent(s): ea51d02

cv2.rectangle attempt

Browse files
Files changed (1) hide show
  1. app.py +17 -11
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
- color = np.random.randint(0, 256, size=3).tolist()
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(image_cv2, contours, -1, color, 2)
 
 
 
 
68
 
69
- return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)
 
 
 
 
 
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[Dict[str, Any]]:
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]