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 ([email protected]) | |
| # -------------------------------------------------------- | |
| import glob | |
| import os | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from torchvision import transforms | |
| from detectron2.data import MetadataCatalog | |
| from utils.visualizer import Visualizer | |
| from xdecoder.language.loss import vl_similarity | |
| from detectron2.utils.colormap import random_color | |
| t = [] | |
| t.append(transforms.Resize((224,224), interpolation=Image.BICUBIC)) | |
| transform_ret = transforms.Compose(t) | |
| t = [] | |
| t.append(transforms.Resize(512, interpolation=Image.BICUBIC)) | |
| transform_grd = transforms.Compose(t) | |
| metadata = MetadataCatalog.get('coco_2017_train_panoptic') | |
| imgs_root = 'images/coco' | |
| img_pths = sorted(glob.glob(os.path.join(imgs_root, '*.jpg'))) | |
| imgs = [Image.open(x).convert('RGB') for x in img_pths] | |
| v_emb = torch.load("v_emb.da") | |
| def region_retrieval(model, image, texts, inpainting_text, *args, **kwargs): | |
| model_novg, model_seg = model | |
| with torch.no_grad(): | |
| # images = [transform_ret(x) for x in imgs] | |
| # images = [np.asarray(x) for x in imgs] | |
| # images = [torch.from_numpy(x.copy()).permute(2,0,1).cuda() for x in images] | |
| # batch_inputs = [{'image': image, 'image_id': 0} for image in images] | |
| # outputs = model_novg.model.evaluate(batch_inputs) | |
| # v_emb = torch.cat([x['captions'][-1:] for x in outputs]) | |
| # v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7) | |
| # torch.save(v_emb, "v_emb.da") | |
| # exit() | |
| texts_ = [[x.strip() if x.strip().endswith('.') else (x.strip() + '.')] for x in texts.split(',')] | |
| model_novg.model.sem_seg_head.predictor.lang_encoder.get_text_embeddings(texts_, is_eval=False, name='caption', prompt=False) | |
| t_emb = getattr(model_novg.model.sem_seg_head.predictor.lang_encoder, '{}_text_embeddings'.format('caption')) | |
| temperature = model_novg.model.sem_seg_head.predictor.lang_encoder.logit_scale | |
| logits = vl_similarity(v_emb, t_emb, temperature) | |
| prob, idx = logits[:,0].softmax(-1).max(0) | |
| image_ori = imgs[idx] | |
| image = transform_grd(image_ori) | |
| width, height = image.size | |
| image = np.asarray(image) | |
| image_ori = np.asarray(image) | |
| images = torch.from_numpy(image.copy()).permute(2,0,1).cuda() | |
| batch_inputs = [{'image': images, 'height': height, 'width': width, 'groundings': {'texts': texts_}}] | |
| model_seg.model.sem_seg_head.predictor.lang_encoder.get_text_embeddings(texts_, is_eval=False, name='caption', prompt=False) | |
| outputs = model_seg.model.evaluate_grounding(batch_inputs, None) | |
| visual = Visualizer(image_ori, metadata=metadata) | |
| grd_masks = (outputs[0]['grounding_mask'] > 0).float().cpu().numpy() | |
| for text, mask in zip([x[0] for x in texts_], grd_masks): | |
| color = random_color(rgb=True, maximum=1).astype(np.int32).tolist() | |
| demo = visual.draw_binary_mask(mask, color=color, text=texts, alpha=0.5) | |
| res = demo.get_image() | |
| torch.cuda.empty_cache() | |
| return Image.fromarray(res), "Selected Image Probability: {:.2f}".format(prob.item()), None |