Martin Tomov commited on
Commit
e886b5d
1 Parent(s): 386f67c

insect region mask

Browse files
Files changed (1) hide show
  1. app.py +5 -11
app.py CHANGED
@@ -54,8 +54,6 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
54
  yellow_background = np.full(image_cv2.shape, (0, 255, 255), dtype=np.uint8)
55
 
56
  for detection in detection_results:
57
- label = detection.label
58
- score = detection.score
59
  box = detection.box
60
  mask = detection.mask
61
 
@@ -63,16 +61,12 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
63
  mask_uint8 = (mask * 255).astype(np.uint8)
64
  contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
65
  cv2.drawContours(yellow_background, contours, -1, (0, 0, 0), cv2.FILLED)
66
-
67
- # Drawing bounding box (border and inside should be yellow)
68
- cv2.rectangle(yellow_background, (box.xmin, box.ymin), (box.xmax, box.ymax), (0, 255, 255), cv2.FILLED) # Fill inside with yellow
69
- cv2.rectangle(yellow_background, (box.xmin, box.ymin), (box.xmax, box.ymax), (0, 255, 255), 2) # Draw border yellow
70
 
71
- if mask is not None:
72
- # Overlay the insect on top of the yellow background
73
- insect_region = image_cv2[box.ymin:box.ymax, box.xmin:box.xmax]
74
- 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])
75
-
76
  return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
77
 
78
  def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult]) -> np.ndarray:
 
54
  yellow_background = np.full(image_cv2.shape, (0, 255, 255), dtype=np.uint8)
55
 
56
  for detection in detection_results:
 
 
57
  box = detection.box
58
  mask = detection.mask
59
 
 
61
  mask_uint8 = (mask * 255).astype(np.uint8)
62
  contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
63
  cv2.drawContours(yellow_background, contours, -1, (0, 0, 0), cv2.FILLED)
 
 
 
 
64
 
65
+ # Extract insect region using mask
66
+ insect_region = cv2.bitwise_and(image_cv2, image_cv2, mask=mask)
67
+
68
+ yellow_background[mask > 0] = insect_region[mask > 0]
69
+
70
  return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
71
 
72
  def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[DetectionResult]) -> np.ndarray: