Spaces:
Runtime error
Runtime error
Thiago Hersan
commited on
Commit
•
4c76544
1
Parent(s):
7e0a636
blends in and out images
Browse files
app.py
CHANGED
@@ -8,20 +8,26 @@ model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-s
|
|
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 |
-
|
13 |
image_total_pixels = mask.shape[0] * mask.shape[1]
|
14 |
label_ids = np.unique(mask)
|
15 |
|
16 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
17 |
id2count = {id: 0 for id in label_ids}
|
18 |
|
19 |
-
for i in range(
|
20 |
-
for j in range(
|
21 |
-
|
22 |
id2count[mask[i, j]] = id2count[mask[i, j]] + 1
|
23 |
|
24 |
-
|
25 |
|
26 |
dataframe = [[
|
27 |
f"{id2label[id]}",
|
@@ -29,7 +35,7 @@ def visualize_instance_seg_mask(mask, id2label):
|
|
29 |
f"{np.sqrt(id2count[id] / image_total_pixels):.2f} m"
|
30 |
] for id in label_ids if 'tree' in id2label[id]]
|
31 |
|
32 |
-
return
|
33 |
|
34 |
|
35 |
def query_image(img):
|
@@ -37,14 +43,14 @@ def query_image(img):
|
|
37 |
inputs = feature_extractor(images=img, return_tensors="pt")
|
38 |
outputs = model(**inputs)
|
39 |
results = feature_extractor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
|
40 |
-
results = visualize_instance_seg_mask(results.numpy(), model.config.id2label)
|
41 |
return results
|
42 |
|
43 |
|
44 |
demo = gr.Interface(
|
45 |
query_image,
|
46 |
-
inputs=[gr.Image()],
|
47 |
-
outputs=["
|
48 |
title="maskformer-swin-large-coco",
|
49 |
allow_flagging="never",
|
50 |
analytics_enabled=None
|
|
|
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(img_in, mask, id2label):
|
12 |
+
img_out = 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 |
+
def get_color(id):
|
17 |
+
id_color = (np.random.randint(0, 2), np.random.randint(0, 4), np.random.randint(0, 256))
|
18 |
+
if 'tree' in id2label[id]:
|
19 |
+
id_color = (0, 140, 0)
|
20 |
+
return id_color
|
21 |
+
|
22 |
+
id2color = {id: get_color(id) for id in label_ids}
|
23 |
id2count = {id: 0 for id in label_ids}
|
24 |
|
25 |
+
for i in range(img_out.shape[0]):
|
26 |
+
for j in range(img_out.shape[1]):
|
27 |
+
img_out[i, j, :] = id2color[mask[i, j]]
|
28 |
id2count[mask[i, j]] = id2count[mask[i, j]] + 1
|
29 |
|
30 |
+
image_res = (0.5 * img_in + 0.5 * img_out) / 255
|
31 |
|
32 |
dataframe = [[
|
33 |
f"{id2label[id]}",
|
|
|
35 |
f"{np.sqrt(id2count[id] / image_total_pixels):.2f} m"
|
36 |
] for id in label_ids if 'tree' in id2label[id]]
|
37 |
|
38 |
+
return image_res, dataframe
|
39 |
|
40 |
|
41 |
def query_image(img):
|
|
|
43 |
inputs = feature_extractor(images=img, return_tensors="pt")
|
44 |
outputs = model(**inputs)
|
45 |
results = feature_extractor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
|
46 |
+
results = visualize_instance_seg_mask(img, results.numpy(), model.config.id2label)
|
47 |
return results
|
48 |
|
49 |
|
50 |
demo = gr.Interface(
|
51 |
query_image,
|
52 |
+
inputs=[gr.Image(label="Input Image")],
|
53 |
+
outputs=[gr.Image(label="Trees"), gr.DataFrame(headers=None, default_value=None, label="Area Info")],
|
54 |
title="maskformer-swin-large-coco",
|
55 |
allow_flagging="never",
|
56 |
analytics_enabled=None
|