Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -22,54 +22,8 @@ def fig2img(fig):
|
|
22 |
return img
|
23 |
|
24 |
def plot(annotations, prompt_process, mask_random_color=True, better_quality=True, retina=True, with_contours=True):
|
25 |
-
|
26 |
-
|
27 |
-
original_h, original_w = ann.orig_shape
|
28 |
-
fig = plt.figure(figsize=(original_w / 100, original_h / 100))
|
29 |
-
plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
|
30 |
-
plt.margins(0, 0)
|
31 |
-
plt.gca().xaxis.set_major_locator(plt.NullLocator())
|
32 |
-
plt.gca().yaxis.set_major_locator(plt.NullLocator())
|
33 |
-
plt.imshow(image)
|
34 |
-
|
35 |
-
if ann.masks is not None:
|
36 |
-
masks = ann.masks.data
|
37 |
-
if better_quality:
|
38 |
-
if isinstance(masks[0], torch.Tensor):
|
39 |
-
masks = np.array(masks.cpu())
|
40 |
-
for i, mask in enumerate(masks):
|
41 |
-
mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
|
42 |
-
masks[i] = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8))
|
43 |
-
|
44 |
-
prompt_process.fast_show_mask(
|
45 |
-
masks,
|
46 |
-
plt.gca(),
|
47 |
-
random_color=mask_random_color,
|
48 |
-
bbox=None,
|
49 |
-
points=None,
|
50 |
-
pointlabel=None,
|
51 |
-
retinamask=retina,
|
52 |
-
target_height=original_h,
|
53 |
-
target_width=original_w,
|
54 |
-
)
|
55 |
-
|
56 |
-
if with_contours:
|
57 |
-
contour_all = []
|
58 |
-
temp = np.zeros((original_h, original_w, 1))
|
59 |
-
for i, mask in enumerate(masks):
|
60 |
-
mask = mask.astype(np.uint8)
|
61 |
-
if not retina:
|
62 |
-
mask = cv2.resize(mask, (original_w, original_h), interpolation=cv2.INTER_NEAREST)
|
63 |
-
contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
|
64 |
-
contour_all.extend(iter(contours))
|
65 |
-
cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2)
|
66 |
-
color = np.array([0 / 255, 0 / 255, 1.0, 0.8])
|
67 |
-
contour_mask = temp / 255 * color.reshape(1, 1, -1)
|
68 |
-
plt.imshow(contour_mask)
|
69 |
-
|
70 |
-
plt.axis("off")
|
71 |
-
plt.close()
|
72 |
-
return fig2img(fig)
|
73 |
|
74 |
def segment_image(input_image, object_name):
|
75 |
try:
|
@@ -78,8 +32,8 @@ def segment_image(input_image, object_name):
|
|
78 |
|
79 |
input_image = Image.fromarray(input_image).convert("RGB")
|
80 |
|
81 |
-
# Run FastSAM model
|
82 |
-
everything_results = model(input_image, retina_masks=True, imgsz=1024, conf=0.
|
83 |
|
84 |
# Prepare a Prompt Process object
|
85 |
prompt_process = FastSAMPrompt(input_image, everything_results, device=device)
|
@@ -90,6 +44,20 @@ def segment_image(input_image, object_name):
|
|
90 |
if not results:
|
91 |
return input_image, f"Could not find '{object_name}' in the image."
|
92 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
93 |
# Plot the results
|
94 |
result_image = plot(annotations=results, prompt_process=prompt_process)
|
95 |
|
|
|
22 |
return img
|
23 |
|
24 |
def plot(annotations, prompt_process, mask_random_color=True, better_quality=True, retina=True, with_contours=True):
|
25 |
+
# ... (keep the existing plot function as is)
|
26 |
+
# This function doesn't need modification for our purposes
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
27 |
|
28 |
def segment_image(input_image, object_name):
|
29 |
try:
|
|
|
32 |
|
33 |
input_image = Image.fromarray(input_image).convert("RGB")
|
34 |
|
35 |
+
# Run FastSAM model with adjusted parameters
|
36 |
+
everything_results = model(input_image, retina_masks=True, imgsz=1024, conf=0.25, iou=0.7)
|
37 |
|
38 |
# Prepare a Prompt Process object
|
39 |
prompt_process = FastSAMPrompt(input_image, everything_results, device=device)
|
|
|
44 |
if not results:
|
45 |
return input_image, f"Could not find '{object_name}' in the image."
|
46 |
|
47 |
+
# Post-process the masks
|
48 |
+
for ann in results:
|
49 |
+
if ann.masks is not None:
|
50 |
+
masks = ann.masks.data
|
51 |
+
if isinstance(masks[0], torch.Tensor):
|
52 |
+
masks = np.array(masks.cpu())
|
53 |
+
for i, mask in enumerate(masks):
|
54 |
+
# Apply more aggressive morphological operations
|
55 |
+
kernel = np.ones((5,5), np.uint8)
|
56 |
+
mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel)
|
57 |
+
mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, kernel)
|
58 |
+
masks[i] = cv2.dilate(mask, kernel, iterations=2)
|
59 |
+
ann.masks.data = masks
|
60 |
+
|
61 |
# Plot the results
|
62 |
result_image = plot(annotations=results, prompt_process=prompt_process)
|
63 |
|