Create reg_ret.py

tasks/reg_ret.py

@@ -0,0 +1,72 @@
+# --------------------------------------------------------
+# 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