Spaces:
Build error
Build error
| # -------------------------------------------------------- | |
| # X-Decoder -- Generalized Decoding for Pixel, Image, and Language | |
| # Copyright (c) 2022 Microsoft | |
| # Licensed under The MIT License [see LICENSE for details] | |
| # Written by Xueyan Zou (xueyan@cs.wisc.edu) | |
| # -------------------------------------------------------- | |
| import os | |
| import cv2 | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from torchvision import transforms | |
| from utils.visualizer import Visualizer | |
| from detectron2.utils.colormap import random_color | |
| from detectron2.data import MetadataCatalog | |
| t = [] | |
| t.append(transforms.Resize(512, interpolation=Image.BICUBIC)) | |
| transform = transforms.Compose(t) | |
| metadata = MetadataCatalog.get('ade20k_panoptic_train') | |
| def open_semseg(model, image, texts, inpainting_text, *args, **kwargs): | |
| stuff_classes = [x.strip() for x in texts.split(',')] | |
| stuff_colors = [random_color(rgb=True, maximum=255).astype(np.int32).tolist() for _ in range(len(stuff_classes))] | |
| stuff_dataset_id_to_contiguous_id = {x:x for x in range(len(stuff_classes))} | |
| MetadataCatalog.get("demo").set( | |
| stuff_colors=stuff_colors, | |
| stuff_classes=stuff_classes, | |
| stuff_dataset_id_to_contiguous_id=stuff_dataset_id_to_contiguous_id, | |
| ) | |
| model.model.sem_seg_head.predictor.lang_encoder.get_text_embeddings(stuff_classes + ["background"], is_eval=True) | |
| metadata = MetadataCatalog.get('demo') | |
| model.model.metadata = metadata | |
| model.model.sem_seg_head.num_classes = len(stuff_classes) | |
| with torch.no_grad(): | |
| image_ori = transform(image) | |
| width = image_ori.size[0] | |
| height = image_ori.size[1] | |
| image = transform(image_ori) | |
| image = np.asarray(image) | |
| images = torch.from_numpy(image.copy()).permute(2,0,1).cuda() | |
| batch_inputs = [{'image': images, 'height': height, 'width': width}] | |
| outputs = model.forward(batch_inputs) | |
| visual = Visualizer(image_ori, metadata=metadata) | |
| sem_seg = outputs[-1]['sem_seg'].max(0)[1] | |
| demo = visual.draw_sem_seg(sem_seg.cpu(), alpha=0.5) # rgb Image | |
| res = demo.get_image() | |
| MetadataCatalog.remove('demo') | |
| torch.cuda.empty_cache() | |
| return Image.fromarray(res), '', None |