Spaces:
Runtime error
Runtime error
import glob | |
import gradio as gr | |
import numpy as np | |
from os import environ | |
from PIL import Image | |
from torchvision import transforms as T | |
from transformers import MaskFormerForInstanceSegmentation, MaskFormerImageProcessor | |
example_images = sorted(glob.glob('examples/map*.jpg')) | |
ade_mean=[0.485, 0.456, 0.406] | |
ade_std=[0.229, 0.224, 0.225] | |
test_transform = T.Compose([ | |
T.ToTensor(), | |
T.Normalize(mean=ade_mean, std=ade_std) | |
]) | |
palette = [ | |
[120, 120, 120], [4, 200, 4], [4, 4, 250], [6, 230, 230], | |
[80, 50, 50], [120, 120, 80], [140, 140, 140], [204, 5, 255] | |
] | |
model_id = f"thiagohersan/maskformer-satellite-trees" | |
vegetation_labels = ["vegetation"] | |
# preprocessor = MaskFormerImageProcessor.from_pretrained(model_id) | |
preprocessor = MaskFormerImageProcessor( | |
do_resize=False, | |
do_normalize=False, | |
do_rescale=False, | |
ignore_index=255, | |
reduce_labels=False | |
) | |
hf_token = environ.get('HFTOKEN') | |
model = MaskFormerForInstanceSegmentation.from_pretrained(model_id, use_auth_token=hf_token) | |
def visualize_instance_seg_mask(img_in, mask, id2label, included_labels): | |
img_out = np.zeros((mask.shape[0], mask.shape[1], 3)) | |
image_total_pixels = mask.shape[0] * mask.shape[1] | |
label_ids = np.unique(mask) | |
id2color = {id: palette[id] for id in label_ids} | |
id2count = {id: 0 for id in label_ids} | |
for i in range(img_out.shape[0]): | |
for j in range(img_out.shape[1]): | |
img_out[i, j, :] = id2color[mask[i, j]] | |
id2count[mask[i, j]] = id2count[mask[i, j]] + 1 | |
image_res = (0.5 * img_in + 0.5 * img_out).astype(np.uint8) | |
dataframe = [[ | |
f"{id2label[id]}", | |
f"{(100 * id2count[id] / image_total_pixels):.2f} %", | |
f"{np.sqrt(id2count[id] / image_total_pixels):.2f} m" | |
] for id in label_ids if id2label[id] in included_labels] | |
if len(dataframe) < 1: | |
dataframe = [[ | |
f"", | |
f"{(0):.2f} %", | |
f"{(0):.2f} m" | |
]] | |
return image_res, dataframe | |
def query_image(image_path): | |
img = np.array(Image.open(image_path)) | |
img_size = (img.shape[0], img.shape[1]) | |
inputs = preprocessor(images=test_transform(img), return_tensors="pt") | |
outputs = model(**inputs) | |
results = preprocessor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0] | |
mask_img, dataframe = visualize_instance_seg_mask(img, results.numpy(), model.config.id2label, vegetation_labels) | |
return mask_img, dataframe | |
def get_system_memory(): | |
memory = psutil.virtual_memory() | |
memory_percent = memory.percent | |
memory_used = memory.used / (1024.0 ** 3) | |
memory_total = memory.total / (1024.0 ** 3) | |
return {"percent": f"{memory_percent}%", "used": f"{memory_used:.3f}GB", "total": f"{memory_total:.3f}GB"} | |
demo = gr.Interface( | |
title="Maskformer Satellite+Trees", | |
description="Using a finetuned version of the [facebook/maskformer-swin-base-ade](https://huggingface.co/facebook/maskformer-swin-base-ade) model (created specifically to work with satellite images) to calculate percentage of pixels in an image that belong to vegetation.", | |
fn=query_image, | |
inputs=[gr.Image(type="filepath", label="Input Image")], | |
outputs=[ | |
gr.Image(label="Vegetation"), | |
gr.DataFrame(label="Info", headers=["Object Label", "Pixel Percent", "Square Length"]) | |
], | |
allow_flagging="never", | |
analytics_enabled=None | |
) | |
demo.launch(show_api=True) | |