sagar007 commited on
Commit
2dd8fe8
·
verified ·
1 Parent(s): 7b13bf4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -50
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
- for ann in annotations:
26
- image = ann.orig_img[..., ::-1] # BGR to RGB
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.4, iou=0.9)
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