Martin Tomov commited on
Commit
5bd1d98
·
verified ·
1 Parent(s): a0de963

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -16
app.py CHANGED
@@ -50,30 +50,25 @@ def annotate(image: Union[Image.Image, np.ndarray], detection_results: List[Dete
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
- # Create an empty yellow background
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
- mask_3_channel = cv2.merge([mask_uint8, mask_uint8, mask_uint8])
62
-
63
- # Invert the mask to get the background part
64
- inverted_mask = cv2.bitwise_not(mask_uint8)
65
- inverted_mask_3_channel = cv2.merge([inverted_mask, inverted_mask, inverted_mask])
66
 
67
  # Extract insect region using mask
68
- insect_region = cv2.bitwise_and(image_cv2, mask_3_channel)
69
- yellow_background_region = cv2.bitwise_and(yellow_background, inverted_mask_3_channel)
70
-
71
- # Combine the insect region and the yellow background region
72
- combined_region = cv2.add(insect_region, yellow_background_region)
73
 
74
- # Apply combined region onto the yellow background
75
- yellow_background = cv2.bitwise_and(yellow_background, inverted_mask_3_channel)
76
- yellow_background = cv2.add(yellow_background, combined_region)
77
 
78
  return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
79
 
@@ -82,7 +77,7 @@ def plot_detections(image: Union[Image.Image, np.ndarray], detections: List[Dete
82
  return annotated_image
83
 
84
  def load_image(image: Union[str, Image.Image]) -> Image.Image:
85
- if isinstance(image, str) and image.startswith("http"):
86
  image = Image.open(requests.get(image, stream=True).raw).convert("RGB")
87
  elif isinstance(image, str):
88
  image = Image.open(image).convert("RGB")
 
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
+ # Create a completely yellow background
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
+ # Drawing bounding box with yellow fill
61
+ cv2.rectangle(yellow_background, (box.xmin, box.ymin), (box.xmax, box.ymax), (0, 255, 255), cv2.FILLED)
62
+
63
  if mask is not None:
64
+ mask = mask.astype(bool)
65
+ mask_3_channel = np.stack([mask, mask, mask], axis=-1)
 
 
 
 
66
 
67
  # Extract insect region using mask
68
+ insect_region = image_cv2 * mask_3_channel
 
 
 
 
69
 
70
+ # Overlay insect region within the bounding box on the yellow background
71
+ yellow_background[mask_3_channel] = insect_region[mask_3_channel]
 
72
 
73
  return cv2.cvtColor(yellow_background, cv2.COLOR_BGR2RGB)
74
 
 
77
  return annotated_image
78
 
79
  def load_image(image: Union[str, Image.Image]) -> Image.Image:
80
+ if isinstance(image, str) and image.startsWith("http"):
81
  image = Image.open(requests.get(image, stream=True).raw).convert("RGB")
82
  elif isinstance(image, str):
83
  image = Image.open(image).convert("RGB")