Spaces:
Build error
Build error
File size: 5,735 Bytes
edebe10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import torch
import numpy as np
import torch.nn.functional as F
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
# from detectron2.structures import BitMasks
from modeling.language.loss import vl_similarity
from utilities.constants import BIOMED_CLASSES
#from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES
# import cv2
# import os
# import glob
# import subprocess
from PIL import Image
import random
t = []
t.append(transforms.Resize((1024, 1024), interpolation=Image.BICUBIC))
transform = transforms.Compose(t)
#metadata = MetadataCatalog.get('coco_2017_train_panoptic')
all_classes = ['background'] + [name.replace('-other','').replace('-merged','')
for name in BIOMED_CLASSES] + ["others"]
# colors_list = [(np.array(color['color'])/255).tolist() for color in COCO_CATEGORIES] + [[1, 1, 1]]
# use color list from matplotlib
import matplotlib.colors as mcolors
colors = dict(mcolors.TABLEAU_COLORS, **mcolors.BASE_COLORS)
colors_list = [list(colors.values())[i] for i in range(16)]
from .output_processing import mask_stats, combine_masks
@torch.no_grad()
def interactive_infer_image(model, image, prompts):
image_resize = transform(image)
width = image.size[0]
height = image.size[1]
image_resize = np.asarray(image_resize)
image = torch.from_numpy(image_resize.copy()).permute(2,0,1).cuda()
data = {"image": image, 'text': prompts, "height": height, "width": width}
# inistalize task
model.model.task_switch['spatial'] = False
model.model.task_switch['visual'] = False
model.model.task_switch['grounding'] = True
model.model.task_switch['audio'] = False
model.model.task_switch['grounding'] = True
batch_inputs = [data]
results,image_size,extra = model.model.evaluate_demo(batch_inputs)
pred_masks = results['pred_masks'][0]
v_emb = results['pred_captions'][0]
t_emb = extra['grounding_class']
t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)
temperature = model.model.sem_seg_head.predictor.lang_encoder.logit_scale
out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
matched_id = out_prob.max(0)[1]
pred_masks_pos = pred_masks[matched_id,:,:]
pred_class = results['pred_logits'][0][matched_id].max(dim=-1)[1]
# interpolate mask to ori size
pred_mask_prob = F.interpolate(pred_masks_pos[None,], (data['height'], data['width']),
mode='bilinear')[0,:,:data['height'],:data['width']].sigmoid().cpu().numpy()
pred_masks_pos = (1*(pred_mask_prob > 0.5)).astype(np.uint8)
return pred_mask_prob
# def interactive_infer_panoptic_biomedseg(model, image, tasks, reftxt=None):
# image_ori = transform(image)
# #mask_ori = image['mask']
# width = image_ori.size[0]
# height = image_ori.size[1]
# image_ori = np.asarray(image_ori)
# visual = Visualizer(image_ori, metadata=metadata)
# images = torch.from_numpy(image_ori.copy()).permute(2,0,1).cuda()
# data = {"image": images, "height": height, "width": width}
# if len(tasks) == 0:
# tasks = ["Panoptic"]
# # inistalize task
# model.model.task_switch['spatial'] = False
# model.model.task_switch['visual'] = False
# model.model.task_switch['grounding'] = False
# model.model.task_switch['audio'] = False
# # check if reftxt is list of strings
# assert isinstance(reftxt, list), f"reftxt should be a list of strings, but got {type(reftxt)}"
# model.model.task_switch['grounding'] = True
# predicts = {}
# for i, txt in enumerate(reftxt):
# data['text'] = txt
# batch_inputs = [data]
# results,image_size,extra = model.model.evaluate_demo(batch_inputs)
# pred_masks = results['pred_masks'][0]
# v_emb = results['pred_captions'][0]
# t_emb = extra['grounding_class']
# t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
# v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)
# temperature = model.model.sem_seg_head.predictor.lang_encoder.logit_scale
# out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
# matched_id = out_prob.max(0)[1]
# pred_masks_pos = pred_masks[matched_id,:,:]
# pred_class = results['pred_logits'][0][matched_id].max(dim=-1)[1]
# # interpolate mask to ori size
# #pred_masks_pos = (F.interpolate(pred_masks_pos[None,], image_size[-2:], mode='bilinear')[0,:,:data['height'],:data['width']] > 0.0).float().cpu().numpy()
# # masks.append(pred_masks_pos[0])
# # mask = pred_masks_pos[0]
# # masks.append(mask)
# # interpolate mask to ori size
# pred_mask_prob = F.interpolate(pred_masks_pos[None,], image_size[-2:], mode='bilinear')[0,:,:data['height'],:data['width']].sigmoid().cpu().numpy()
# #pred_masks_pos = 1*(pred_mask_prob > 0.5)
# predicts[txt] = pred_mask_prob[0]
# masks = combine_masks(predicts)
# predict_mask_stats = {}
# print(masks.keys())
# for i, txt in enumerate(masks):
# mask = masks[txt]
# demo = visual.draw_binary_mask(mask, color=colors_list[i], text=txt)
# predict_mask_stats[txt] = mask_stats((predicts[txt]*255), image_ori)
# res = demo.get_image()
# torch.cuda.empty_cache()
# # return Image.fromarray(res), stroke_inimg, stroke_refimg
# return Image.fromarray(res), None, predict_mask_stats
|