diff --git a/projects/demo.py b/projects/demo.py new file mode 100644 index 0000000..9a038d3 --- /dev/null +++ b/projects/demo.py @@ -0,0 +1,272 @@ +import detectron2 +import numpy as np +import cv2 +import torch +from os import path +from detectron2.config import get_cfg +from GLEE.glee.models.glee_model import GLEE_Model +from GLEE.glee.config_deeplab import add_deeplab_config +from GLEE.glee.config import add_glee_config +import torch.nn.functional as F +import torchvision +import math +from scipy.optimize import linear_sum_assignment +import argparse +from PIL import Image +import os + +def get_args_parser(): + parser = argparse.ArgumentParser('Set transformer detector', add_help=False) + parser.add_argument('--version', type=str, default='Lite', help='select model version from [Lite,Plus,Pro]') + parser.add_argument('--input_image', type=str, default='./Examples/000000001000.jpg', help='path to image') + parser.add_argument('--output', type=str, default='./outputs', help='path to save detection results') + parser.add_argument('--task', type=str, default='detection', help='mode: detection/grounding') + parser.add_argument('--text', type=str, default='person,bicycle,car,motorcycle,airplane', help='category list split by ,\ or a sentence') + parser.add_argument('--topk', type=int, default=10) + parser.add_argument('--sim_thres', type=float, default=0.1, help='Similarity Threshold') + return parser + + +def box_cxcywh_to_xyxy(x): + x_c, y_c, w, h = x.unbind(-1) + b = [(x_c - 0.5 * w), (y_c - 0.5 * h), + (x_c + 0.5 * w), (y_c + 0.5 * h)] + return torch.stack(b, dim=-1) + + +def LSJ_box_postprocess( out_bbox, padding_size, crop_size, img_h, img_w): + # postprocess box height and width + boxes = box_cxcywh_to_xyxy(out_bbox) + lsj_sclae = torch.tensor([padding_size[1], padding_size[0], padding_size[1], padding_size[0]]).to(out_bbox) + crop_scale = torch.tensor([crop_size[1], crop_size[0], crop_size[1], crop_size[0]]).to(out_bbox) + boxes = boxes * lsj_sclae + boxes = boxes / crop_scale + boxes = torch.clamp(boxes,0,1) + + scale_fct = torch.tensor([img_w, img_h, img_w, img_h]) + scale_fct = scale_fct.to(out_bbox) + boxes = boxes * scale_fct + return boxes + +COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125], + [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933], + [0.494, 0.000, 0.556], [0.494, 0.000, 0.000], [0.000, 0.745, 0.000], + [0.700, 0.300, 0.600],[0.000, 0.447, 0.741], [0.850, 0.325, 0.098]] + + +def main(args): + + + print(f"Is CUDA available: {torch.cuda.is_available()}") + # True + if torch.cuda.is_available(): + print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}") + # Tesla T4 + + coco_class_name = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] + YTBVISOVIS_class_name = ['lizard', 'cat', 'horse', 'eagle', 'frog', 'Horse', 'monkey', 'bear', 'parrot', 'giant_panda', 'truck', 'zebra', 'rabbit', 'skateboard', 'tiger', 'shark', 'Person', 'Poultry', 'Zebra', 'Airplane', 'elephant', 'Elephant', 'Turtle', 'snake', 'train', 'Dog', 'snowboard', 'airplane', 'Lizard', 'dog', 'Cat', 'earless_seal', 'boat', 'Tiger', 'motorbike', 'duck', 'fox', 'Monkey', 'Bird', 'Bear', 'tennis_racket', 'Rabbit', 'Giraffe', 'Motorcycle', 'fish', 'Boat', 'deer', 'ape', 'Bicycle', 'Parrot', 'Cow', 'turtle', 'mouse', 'owl', 'Fish', 'surfboard', 'Giant_panda', 'Sheep', 'hand', 'Vehical', 'sedan', 'leopard', 'person', 'giraffe', 'cow'] + class_agnostic_name = ['object'] + + if torch.cuda.is_available(): + print('use cuda') + device = 'cuda' + else: + print('use cpu') + device='cpu' + + if 'Lite' in args.version: + cfg_r50 = get_cfg() + add_deeplab_config(cfg_r50) + add_glee_config(cfg_r50) + conf_files_r50 = 'GLEE/configs/R50.yaml' + checkpoints_r50 = torch.load('GLEE_DEMO_MODEL_ZOO/GLEE_R50_Scaleup10m.pth') + cfg_r50.merge_from_file(conf_files_r50) + GLEEmodel = GLEE_Model(cfg_r50, None, device, None, True).to(device) + GLEEmodel.load_state_dict(checkpoints_r50, strict=False) + GLEEmodel.eval() + inference_type = 'resize_shot' # or LSJ + elif 'Plus' in args.version: + cfg_swin = get_cfg() + add_deeplab_config(cfg_swin) + add_glee_config(cfg_swin) + conf_files_swin = 'GLEE/configs/SwinL.yaml' + checkpoints_swin = torch.load('GLEE_DEMO_MODEL_ZOO/GLEE_SwinL_Scaleup10m.pth') + cfg_swin.merge_from_file(conf_files_swin) + GLEEmodel = GLEE_Model(cfg_swin, None, device, None, True).to(device) + GLEEmodel.load_state_dict(checkpoints_swin, strict=False) + GLEEmodel.eval() + inference_type = 'resize_shot' # or LSJ + elif 'Pro' in args.version: + cfg_eva02 = get_cfg() + add_deeplab_config(cfg_eva02) + add_glee_config(cfg_eva02) + conf_files_eva02 = 'GLEE/configs/EVA02.yaml' + checkpoints_eva = torch.load('GLEE_DEMO_MODEL_ZOO/GLEE_EVA02_Scaleup10m.pth') + cfg_eva02.merge_from_file(conf_files_eva02) + GLEEmodel = GLEE_Model(cfg_eva02, None, device, None, True).to(device) + GLEEmodel.load_state_dict(checkpoints_eva, strict=False) + GLEEmodel.eval() + inference_type = 'LSJ' + else: + assert False, 'model version not defined!' + + + + + pixel_mean = torch.Tensor( [123.675, 116.28, 103.53]).to(device).view(3, 1, 1) + pixel_std = torch.Tensor([58.395, 57.12, 57.375]).to(device).view(3, 1, 1) + normalizer = lambda x: (x - pixel_mean) / pixel_std + inference_size = 800 + + size_divisibility = 32 + + FONT_SCALE = 1.5e-3 + THICKNESS_SCALE = 1e-3 + TEXT_Y_OFFSET_SCALE = 1e-2 + + + if inference_type != 'LSJ': + resizer = torchvision.transforms.Resize(inference_size,antialias=True) + else: + resizer = torchvision.transforms.Resize(size = 1535, max_size=1536, antialias=True) + + + inputimage = np.array(Image.open(args.input_image)) + + + + ori_image = torch.as_tensor(np.ascontiguousarray( inputimage.transpose(2, 0, 1))) + ori_image = normalizer(ori_image.to(device))[None,] + _,_, ori_height, ori_width = ori_image.shape + + if inference_type == 'LSJ': + resize_image = resizer(ori_image) + image_size = torch.as_tensor((resize_image.shape[-2],resize_image.shape[-1])) + re_size = resize_image.shape[-2:] + infer_image = torch.zeros(1,3,1536,1536).to(ori_image) + infer_image[:,:,:image_size[0],:image_size[1]] = resize_image + padding_size = (1536,1536) + else: + resize_image = resizer(ori_image) + image_size = torch.as_tensor((resize_image.shape[-2],resize_image.shape[-1])) + re_size = resize_image.shape[-2:] + if size_divisibility > 1: + stride = size_divisibility + # the last two dims are H,W, both subject to divisibility requirement + padding_size = ((image_size + (stride - 1)).div(stride, rounding_mode="floor") * stride).tolist() + infer_image = torch.zeros(1,3,padding_size[0],padding_size[1]).to(resize_image) + infer_image[0,:,:image_size[0],:image_size[1]] = resize_image + # reversed_image = infer_image*pixel_std + pixel_mean + # reversed_image = torch.clip(reversed_image,min=0,max=255) + # reversed_image = reversed_image[0].permute(1,2,0) + # reversed_image = reversed_image.int().cpu().numpy().copy() + # cv2.imwrite('test.png',reversed_image[:,:,::-1]) + + results_select=['box','name','score'] # or ['box','mask'] #选择要可视化的部分 + topK_instance = args.topk + threshold_select = args.sim_thres + + if args.task == 'detection': + batch_category_name = args.text.split(',') + prompt_list = [] + task="coco" + elif args.task == 'grounding': + batch_category_name = [] + prompt_list = {'grounding':[args.text]} + task="grounding" + else: + assert False, 'task not defined!' + + with torch.no_grad(): + (outputs,_) = GLEEmodel(infer_image, prompt_list, task=task, batch_name_list=batch_category_name, is_train=False) + + mask_pred = outputs['pred_masks'][0] + mask_cls = outputs['pred_logits'][0] + boxes_pred = outputs['pred_boxes'][0] + + scores = mask_cls.sigmoid().max(-1)[0] + scores_per_image, topk_indices = scores.topk(topK_instance, sorted=True) + + valid = scores_per_image>threshold_select + topk_indices = topk_indices[valid] + scores_per_image = scores_per_image[valid] + + pred_class = mask_cls[topk_indices].max(-1)[1].tolist() + pred_boxes = boxes_pred[topk_indices] + + + boxes = LSJ_box_postprocess(pred_boxes,padding_size,re_size, ori_height,ori_width) + mask_pred = mask_pred[topk_indices] + assert len(mask_pred)>0 ,'not enough object to visualize, turn thres bigger' + pred_masks = F.interpolate( mask_pred[None,], size=(padding_size[0], padding_size[1]), mode="bilinear", align_corners=False ) + pred_masks = pred_masks[:,:,:re_size[0],:re_size[1]] + pred_masks = F.interpolate( pred_masks, size=(ori_height,ori_width), mode="bilinear", align_corners=False ) + pred_masks = (pred_masks>0).detach().cpu().numpy()[0] + + if 'mask' in results_select: + mask_image_mix_ration=0.5 + zero_mask = np.zeros_like(inputimage) + for nn, mask in enumerate(pred_masks): + # mask = mask.numpy() + mask = mask.reshape(mask.shape[0], mask.shape[1], 1) + + lar = np.concatenate((mask*COLORS[nn%12][2], mask*COLORS[nn%12][1], mask*COLORS[nn%12][0]), axis = 2) + zero_mask = zero_mask+ lar + + + lar_valid = zero_mask>0 + masked_image = lar_valid*inputimage + img_n = masked_image*mask_image_mix_ration + np.clip(zero_mask,0,1)*255*(1-mask_image_mix_ration) + max_p = img_n.max() + img_n = 255*img_n/max_p + ret = (~lar_valid*inputimage)*mask_image_mix_ration + img_n + ret = ret.astype('uint8') + else: + ret = inputimage + + if 'box' in results_select: + + line_width = max(ret.shape) /200 + + for nn,(classid, box) in enumerate(zip(pred_class,boxes)): + x1,y1,x2,y2 = box.long().tolist() + RGB = (COLORS[nn%12][2]*255,COLORS[nn%12][1]*255,COLORS[nn%12][0]*255) + cv2.rectangle(ret, (x1,y1), (x2,y2), RGB, math.ceil(line_width) ) + if args.task == 'detection' : + label = '' + if 'name' in results_select: + label += batch_category_name[classid] + if 'score' in results_select: + label += str(scores_per_image[nn].item())[:3] + + if len(label)==0: + continue + height, width, _ = ret.shape + FONT = cv2.FONT_HERSHEY_COMPLEX + label_width, label_height = cv2.getTextSize(label, FONT, min(width, height) * FONT_SCALE, math.ceil(min(width, height) * THICKNESS_SCALE))[0] + + cv2.rectangle(ret, (x1,y1), (x1+label_width,(y1 -label_height) - int(height * TEXT_Y_OFFSET_SCALE)), RGB, -1) + + cv2.putText( + ret, + label, + (x1, y1 - int(height * TEXT_Y_OFFSET_SCALE)), + fontFace=FONT, + fontScale=min(width, height) * FONT_SCALE, + thickness=math.ceil(min(width, height) * THICKNESS_SCALE), + color=(255,255,255), + ) + + ret = ret.astype('uint8') + if not os.path.exists(args.output): + os.makedirs(args.output) + Image.fromarray(ret).save(os.path.join(args.output, args.input_image.split('/')[-1])) + # cv2.imwrite( os.path.join(args.output, args.input_image.split('/')[-1]),ret ) + + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser('image path check script', parents=[get_args_parser()]) + args = parser.parse_args() + main(args) \ No newline at end of file