Martin Tomov commited on
Commit
02b116a
Β·
verified Β·
1 Parent(s): 5977b5c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -9
app.py CHANGED
@@ -54,19 +54,12 @@ 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
- box = detection.box
58
  mask = detection.mask
59
 
60
  if mask is not None:
61
  mask_uint8 = (mask * 255).astype(np.uint8)
62
- mask_inv = cv2.bitwise_not(mask_uint8)
63
-
64
- # Extract insect using mask
65
  insect_region = cv2.bitwise_and(image_cv2, image_cv2, mask=mask_uint8)
66
- yellow_background_masked = cv2.bitwise_and(yellow_background, yellow_background, mask=mask_inv)
67
- insect_combined = cv2.add(insect_region, yellow_background_masked)
68
-
69
- yellow_background = cv2.add(yellow_background, insect_combined)
70
 
71
  return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
72
 
@@ -218,4 +211,4 @@ gr.Interface(
218
  outputs=[gr.Image(type="numpy"), gr.Textbox()],
219
  title="InsectSAM 🐞",
220
  examples=examples
221
- ).launch()
 
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_uint8 = (mask * 255).astype(np.uint8)
 
 
 
61
  insect_region = cv2.bitwise_and(image_cv2, image_cv2, mask=mask_uint8)
62
+ yellow_background = np.where(mask[..., None], insect_region, yellow_background)
 
 
 
63
 
64
  return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
65
 
 
211
  outputs=[gr.Image(type="numpy"), gr.Textbox()],
212
  title="InsectSAM 🐞",
213
  examples=examples
214
+ ).launch()