gatesla commited on
Commit
8b5e39e
1 Parent(s): a134e44

Got it to successfully create individual bounding boxes for a whole mask

Browse files
Files changed (2) hide show
  1. app.py +5 -5
  2. understand.py +74 -0
app.py CHANGED
@@ -6,7 +6,7 @@ import torch
6
  import pathlib
7
  from PIL import Image
8
 
9
- from transformers import DetrFeatureExtractor, DetrForSegmentation, MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
10
  from transformers.models.detr.feature_extraction_detr import rgb_to_id
11
 
12
 
@@ -73,15 +73,15 @@ def segment_images(model_name,url_input,image_input,threshold):
73
  pass
74
  elif "maskformer" in model_name.lower():
75
  # Load the processor and model
76
- processor = MaskFormerFeatureExtractor.from_pretrained(model_name)
77
- print(type(processor))
78
  model = MaskFormerForInstanceSegmentation.from_pretrained(model_name)
79
 
80
  inputs = processor(images=image, return_tensors="pt")
81
 
82
  outputs = model(**inputs)
83
- print(type(outputs))
84
- print(outputs)
85
  pass
86
  else:
87
  raise NameError("Model is not implemented")
 
6
  import pathlib
7
  from PIL import Image
8
 
9
+ from transformers import DetrFeatureExtractor, DetrForSegmentation, MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
10
  from transformers.models.detr.feature_extraction_detr import rgb_to_id
11
 
12
 
 
73
  pass
74
  elif "maskformer" in model_name.lower():
75
  # Load the processor and model
76
+ processor = MaskFormerImageProcessor.from_pretrained(model_name)
77
+ # print(type(processor))
78
  model = MaskFormerForInstanceSegmentation.from_pretrained(model_name)
79
 
80
  inputs = processor(images=image, return_tensors="pt")
81
 
82
  outputs = model(**inputs)
83
+ results = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
84
+
85
  pass
86
  else:
87
  raise NameError("Model is not implemented")
understand.py CHANGED
@@ -63,6 +63,67 @@ results = processor.post_process_panoptic_segmentation(outputs, target_sizes=[im
63
  # type(results["segmentation"]) --> <class 'torch.Tensor'>
64
 
65
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66
 
67
  # From Tutorial (Box 79)
68
  # def get_mask(segment_idx):
@@ -129,4 +190,17 @@ array([[False, False, False, ..., False, False, False],
129
  >>> results["segments_info"][0]
130
  {'id': 1, 'label_id': 25, 'was_fused': False, 'score': 0.998022}
131
  >>>
 
 
 
 
 
 
 
 
 
 
 
 
 
132
  """
 
63
  # type(results["segmentation"]) --> <class 'torch.Tensor'>
64
 
65
 
66
+ def show_mask_for_number(map_to_use, label_id):
67
+
68
+ if torch.cuda.is_available():
69
+ mask = (map_to_use.cpu().numpy() == label_id)
70
+ else:
71
+ mask = (map_to_use.numpy() == label_id)
72
+
73
+ visual_mask = (mask* 255).astype(np.uint8)
74
+ visual_mask = Image.fromarray(visual_mask)
75
+ plt.imshow(visual_mask)
76
+ plt.show()
77
+
78
+
79
+ def get_coordinates_for_bb_simple(map_to_use, label_id):
80
+ if torch.cuda.is_available():
81
+ mask = (map_to_use.cpu().numpy() == label_id)
82
+ else:
83
+ mask = (map_to_use.numpy() == label_id)
84
+
85
+ x, y = np.where(mask==True)
86
+ x_max, x_min = max(x), min(x)
87
+ y_max, y_min = max(y), min(y)
88
+ return (x_min, y_min), (x_max, y_max)
89
+
90
+ def make_simple_box(left_top, right_bottom, map_size):
91
+ full_mask = np.full(map_size, False)
92
+ left_x, top_y = left_top
93
+ right_x, bottom_y = right_bottom
94
+ full_mask[left_x:right_x, top_y] = True
95
+ full_mask[left_x:right_x, bottom_y] = True
96
+ full_mask[left_x, top_y:bottom_y] = True
97
+ full_mask[right_x, top_y:bottom_y] = True
98
+
99
+ visual_mask = (full_mask* 255).astype(np.uint8)
100
+ visual_mask = Image.fromarray(visual_mask)
101
+ plt.imshow(visual_mask)
102
+ plt.show()
103
+
104
+
105
+ def test(map_to_use, label_id):
106
+ if torch.cuda.is_available():
107
+ mask = (map_to_use.cpu().numpy() == label_id)
108
+ else:
109
+ mask = (map_to_use.numpy() == label_id)
110
+
111
+
112
+ lt, rb = get_coordinates_for_bb_simple(map_to_use, label_id)
113
+ left_x, top_y = lt
114
+ right_x, bottom_y = rb
115
+
116
+ mask[left_x:right_x, top_y] = .5
117
+ mask[left_x:right_x, bottom_y] = .5
118
+ mask[left_x, top_y:bottom_y] = .5
119
+ mask[right_x, top_y:bottom_y] = .5
120
+
121
+ visual_mask = (mask* 255).astype(np.uint8)
122
+ visual_mask = Image.fromarray(visual_mask)
123
+ plt.imshow(visual_mask)
124
+ plt.show()
125
+
126
+
127
 
128
  # From Tutorial (Box 79)
129
  # def get_mask(segment_idx):
 
190
  >>> results["segments_info"][0]
191
  {'id': 1, 'label_id': 25, 'was_fused': False, 'score': 0.998022}
192
  >>>
193
+ """
194
+
195
+ """
196
+ >>> np.where(mask==True)
197
+ (array([300, 300, 300, ..., 392, 392, 392]), array([452, 453, 454, ..., 473, 474, 475]))
198
+ >>> max(np.where(mask==True)[0])
199
+ 392
200
+ >>> min(np.where(mask==True)[0])
201
+ 300
202
+ >>> max(np.where(mask==True)[1])
203
+ 538
204
+ >>> min(np.where(mask==True)[1])
205
+ 399
206
  """