| import gradio as gr | |
| import glob | |
| import torch | |
| import pickle | |
| from PIL import Image, ImageDraw | |
| import numpy as np | |
| from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation | |
| import numpy as np | |
| from scipy.ndimage import center_of_mass | |
| def combine_ims(im1, im2, val=128): | |
| p = Image.new("L", im1.size, val) | |
| im = Image.composite(im1, im2, p) | |
| return im | |
| def get_class_centers(segmentation_mask, class_dict): | |
| segmentation_mask = segmentation_mask.numpy() + 1 | |
| class_centers = {} | |
| for class_index, _ in class_dict.items(): | |
| class_mask = (segmentation_mask == class_index).astype(int) | |
| center_of_mass_list = center_of_mass(class_mask) | |
| class_centers[class_index] = center_of_mass_list | |
| class_centers = {k:list(map(int, v)) for k,v in class_centers.items() if not np.isnan(sum(v))} | |
| return class_centers | |
| def visualize_mask(predicted_semantic_map, class_ids, class_colors): | |
| h, w = predicted_semantic_map.shape | |
| color_indexes = np.zeros((h, w), dtype=np.uint8) | |
| color_indexes[:] = predicted_semantic_map.numpy() | |
| color_indexes = color_indexes.flatten() | |
| colors = class_colors[class_ids[color_indexes]] | |
| output = colors.reshape(h, w, 3).astype(np.uint8) | |
| image_mask = Image.fromarray(output) | |
| return image_mask | |
| def get_out_image(image, predicted_semantic_map): | |
| class_centers = get_class_centers(predicted_semantic_map, class_dict) | |
| mask = visualize_mask(predicted_semantic_map, class_ids, class_colors) | |
| image_mask = combine_ims(image, mask, val=128) | |
| draw = ImageDraw.Draw(image_mask) | |
| for id, (y, x) in class_centers.items(): | |
| draw.text((x, y), str(class_names[id-1]), fill='black') | |
| return image_mask | |
| def gradio_process(image): | |
| inputs = processor(images=image, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| predicted_semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0] | |
| out_image = get_out_image(image, predicted_semantic_map) | |
| return out_image | |
| with open('ade20k_classes.pickle', 'rb') as f: | |
| class_names, class_ids, class_colors = pickle.load(f) | |
| class_names, class_ids, class_colors = np.array(class_names), np.array(class_ids), np.array(class_colors) | |
| class_dict = dict(zip(class_ids, class_names)) | |
| device = torch.device("cpu") | |
| processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-large-ade-semantic") | |
| model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-large-ade-semantic").to(device) | |
| model.eval() | |
| demo = gr.Interface( | |
| gradio_process, | |
| inputs=gr.inputs.Image(type="pil"), | |
| outputs=gr.outputs.Image(type="pil"), | |
| title="Semantic Interior Segmentation", | |
| examples=glob.glob('./examples/*.jpg'), | |
| allow_flagging="never", | |
| ) | |
| demo.launch() |