Martin Tomov commited on
Commit
014d0b1
·
verified ·
1 Parent(s): 8a3f59a

millionth experiment with bbox

Browse files
Files changed (1) hide show
  1. app.py +11 -16
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
- 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 +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[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]
@@ -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 = {