Spaces:
Running
on
Zero
Running
on
Zero
Martin Tomov
commited on
Update app.py
Browse files
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
|
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 |
-
|
61 |
-
mask_3_channel =
|
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 =
|
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 |
-
#
|
75 |
-
yellow_background =
|
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.
|
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")
|