Spaces:
Runtime error
Runtime error
Thiago Hersan
commited on
Commit
•
480594f
1
Parent(s):
a49b93f
with text/semantic labels
Browse files
app.py
CHANGED
@@ -3,39 +3,38 @@ import numpy as np
|
|
3 |
from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
|
4 |
|
5 |
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
|
10 |
|
11 |
-
|
12 |
-
def visualize_instance_seg_mask(mask):
|
13 |
image = np.zeros((mask.shape[0], mask.shape[1], 3))
|
14 |
image_total_pixels = mask.shape[0] * mask.shape[1]
|
15 |
-
|
16 |
|
17 |
-
|
18 |
-
|
19 |
|
20 |
for i in range(image.shape[0]):
|
21 |
for j in range(image.shape[1]):
|
22 |
-
image[i, j, :] =
|
23 |
-
|
24 |
|
25 |
image = image / 255
|
26 |
|
27 |
-
for
|
28 |
-
|
29 |
|
30 |
return image
|
31 |
|
32 |
|
33 |
def query_image(img):
|
34 |
img_size = (img.shape[0], img.shape[1])
|
35 |
-
inputs =
|
36 |
outputs = model(**inputs)
|
37 |
-
results =
|
38 |
-
results = visualize_instance_seg_mask(results.numpy())
|
39 |
return results
|
40 |
|
41 |
|
|
|
3 |
from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
|
4 |
|
5 |
|
6 |
+
feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-coco")
|
7 |
+
model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-coco")
|
8 |
+
# feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
|
9 |
+
# model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
|
10 |
|
11 |
+
def visualize_instance_seg_mask(mask, id2label):
|
|
|
12 |
image = np.zeros((mask.shape[0], mask.shape[1], 3))
|
13 |
image_total_pixels = mask.shape[0] * mask.shape[1]
|
14 |
+
label_ids = np.unique(mask)
|
15 |
|
16 |
+
id2color = {id: (np.random.randint(0, 2), np.random.randint(0, 256), np.random.randint(0, 256)) for id in label_ids}
|
17 |
+
id2count = {id: 0 for id in label_ids}
|
18 |
|
19 |
for i in range(image.shape[0]):
|
20 |
for j in range(image.shape[1]):
|
21 |
+
image[i, j, :] = id2color[mask[i, j]]
|
22 |
+
id2count[mask[i, j]] = id2count[mask[i, j]] + 1
|
23 |
|
24 |
image = image / 255
|
25 |
|
26 |
+
label2count = {id2label[id]: id2count[id] / image_total_pixels for id in label_ids}
|
27 |
+
print(label2count)
|
28 |
|
29 |
return image
|
30 |
|
31 |
|
32 |
def query_image(img):
|
33 |
img_size = (img.shape[0], img.shape[1])
|
34 |
+
inputs = feature_extractor(images=img, return_tensors="pt")
|
35 |
outputs = model(**inputs)
|
36 |
+
results = feature_extractor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
|
37 |
+
results = visualize_instance_seg_mask(results.numpy(), model.config.id2label)
|
38 |
return results
|
39 |
|
40 |
|