GLEE/projects/demo.py

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)