Got it to successfully create individual bounding boxes for a whole mask
Browse files- app.py +5 -5
- 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,
|
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 =
|
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 |
-
|
84 |
-
|
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 |
"""
|