Spaces:
Runtime error
Runtime 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 |