sagar007 commited on
Commit
72f4c5c
·
verified ·
1 Parent(s): 2dd8fe8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -2
app.py CHANGED
@@ -22,8 +22,54 @@ 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
- # ... (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:
 
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: