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 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_panoseg(model, image, texts, inpainting_text, *args, **kwargs): | |
| stuff_classes = [x.strip() for x in texts.split(';')[0].replace('stuff:','').split(',')] | |
| thing_classes = [x.strip() for x in texts.split(';')[1].replace('thing:','').split(',')] | |
| thing_colors = [random_color(rgb=True, maximum=255).astype(np.int32).tolist() for _ in range(len(thing_classes))] | |
| stuff_colors = [random_color(rgb=True, maximum=255).astype(np.int32).tolist() for _ in range(len(stuff_classes))] | |
| thing_dataset_id_to_contiguous_id = {x:x for x in range(len(thing_classes))} | |
| stuff_dataset_id_to_contiguous_id = {x+len(thing_classes):x for x in range(len(stuff_classes))} | |
| MetadataCatalog.get("demo").set( | |
| thing_colors=thing_colors, | |
| thing_classes=thing_classes, | |
| thing_dataset_id_to_contiguous_id=thing_dataset_id_to_contiguous_id, | |
| 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(thing_classes + stuff_classes + ["background"], is_eval=True) | |
| metadata = MetadataCatalog.get('demo') | |
| model.model.metadata = metadata | |
| model.model.sem_seg_head.num_classes = len(thing_classes + 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) | |
| pano_seg = outputs[-1]['panoptic_seg'][0] | |
| pano_seg_info = outputs[-1]['panoptic_seg'][1] | |
| for i in range(len(pano_seg_info)): | |
| if pano_seg_info[i]['category_id'] in metadata.thing_dataset_id_to_contiguous_id.keys(): | |
| pano_seg_info[i]['category_id'] = metadata.thing_dataset_id_to_contiguous_id[pano_seg_info[i]['category_id']] | |
| else: | |
| pano_seg_info[i]['isthing'] = False | |
| pano_seg_info[i]['category_id'] = metadata.stuff_dataset_id_to_contiguous_id[pano_seg_info[i]['category_id']] | |
| demo = visual.draw_panoptic_seg(pano_seg.cpu(), pano_seg_info) # rgb Image | |
| res = demo.get_image() | |
| MetadataCatalog.remove('demo') | |
| torch.cuda.empty_cache() | |
| return Image.fromarray(res), '', None |