app_util.py

import argparse
import io

import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision.transforms as T
from PIL import Image

from models.blip2_decoder import BLIP2Decoder
from models.deformable_detr.backbone import build_backbone
from models.contextdet_blip2 import ContextDET
from models.post_process import CondNMSPostProcess
from models.transformer import build_ov_transformer
from util.misc import nested_tensor_from_tensor_list


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', type=str, default='cpu')

    parser.add_argument('--lr_backbone_names', default=["backbone.0"], type=str, nargs='+')
    parser.add_argument('--lr_backbone', default=2e-5, type=float)

    parser.add_argument('--with_box_refine', default=True, action='store_false')
    parser.add_argument('--two_stage', default=True, action='store_false')

    # * Backbone
    parser.add_argument('--backbone', default='resnet50', type=str,
                        help="Name of the convolutional backbone to use")
    parser.add_argument('--dilation', action='store_true',
                        help="If true, we replace stride with dilation in the last convolutional block (DC5)")
    parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'),
                        help="Type of positional embedding to use on top of the image features")
    parser.add_argument('--position_embedding_scale', default=2 * np.pi, type=float,
                        help="position / size * scale")
    parser.add_argument('--num_feature_levels', default=5, type=int, help='number of feature levels')

    # * Transformer
    parser.add_argument('--enc_layers', default=6, type=int,
                        help="Number of encoding layers in the transformer")
    parser.add_argument('--dec_layers', default=6, type=int,
                        help="Number of decoding layers in the transformer")
    parser.add_argument('--dim_feedforward', default=2048, type=int,
                        help="Intermediate size of the feedforward layers in the transformer blocks")
    parser.add_argument('--hidden_dim', default=256, type=int,
                        help="Size of the embeddings (dimension of the transformer)")
    parser.add_argument('--dropout', default=0.0, type=float,
                        help="Dropout applied in the transformer")
    parser.add_argument('--nheads', default=8, type=int,
                        help="Number of attention heads inside the transformer's attentions")
    parser.add_argument('--num_queries', default=900, type=int,
                        help="Number of query slots")
    parser.add_argument('--dec_n_points', default=4, type=int)
    parser.add_argument('--enc_n_points', default=4, type=int)

    # * Segmentation
    parser.add_argument('--masks', action='store_true',
                        help="Train segmentation head if the flag is provided")

    parser.add_argument('--assign_first_stage', default=True, action='store_false')
    parser.add_argument('--assign_second_stage', default=True, action='store_false')

    parser.add_argument('--name', default='ov')
    parser.add_argument('--llm_name', default='bert-base-cased')

    parser.add_argument('--resume', default='', type=str)
    return parser.parse_args()


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]
]


def fig2img(fig):
    buf = io.BytesIO()
    fig.savefig(buf)
    buf.seek(0)
    img = Image.open(buf)
    return img


def visualize_prediction(pil_img, output_dict, threshold=0.7):
    keep = output_dict["scores"] > threshold
    boxes = output_dict["boxes"][keep].tolist()
    scores = output_dict["scores"][keep].tolist()
    keep_list = keep.nonzero().squeeze(1).numpy().tolist()
    labels = [output_dict["names"][i] for i in keep_list]

    plt.figure(figsize=(12.8, 8))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100
    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=color, linewidth=3))
        ax.text(xmin, ymin, f"{label}: {score:0.2f}", fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
    plt.axis("off")
    return fig2img(plt.gcf())


class ContextDetDemo():
    def __init__(self, resume):
        self.transform = T.Compose([
            T.Resize(640),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        args = parse_args()

        args.llm_name = 'caption_coco_opt2.7b'
        args.resume = resume

        args.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        num_classes = 2
        device = torch.device(args.device)

        backbone = build_backbone(args)
        transformer = build_ov_transformer(args)
        llm_decoder = BLIP2Decoder(args.llm_name)
        model = ContextDET(
            backbone,
            transformer,
            num_classes=num_classes,
            num_queries=args.num_queries,
            num_feature_levels=args.num_feature_levels,
            aux_loss=False,
            with_box_refine=args.with_box_refine,
            two_stage=args.two_stage,
            llm_decoder=llm_decoder,
        )
        model = model.to(device)

        checkpoint = torch.load(args.resume, map_location='cpu')
        missing_keys, unexpected_keys = model.load_state_dict(checkpoint['model'], strict=False)
        if len(missing_keys) > 0:
            print('Missing Keys: {}'.format(missing_keys))
        if len(unexpected_keys) > 0:
            print('Unexpected Keys: {}'.format(unexpected_keys))

        postprocessor = CondNMSPostProcess(args.num_queries)

        self.model = model
        self.model.eval()
        self.postprocessor = postprocessor

    def forward(self, image, text, task_button, history, threshold=0.3):
        samples = self.transform(image).unsqueeze(0)
        samples = nested_tensor_from_tensor_list(samples)
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        samples = samples.to(device)
        vis = self.model.llm_decoder.vis_processors

        if task_button == "Question Answering":
            text = f"{text} Answer:"
            history.append(text)
            # prompt = " ".join(history)
            prompt = text
        elif task_button == "Captioning":
            prompt = "A photo of"
        else:
            prompt = text

        blip2_samples = {
            'image': vis['eval'](image)[None, :].to(device),
            'prompt': [prompt],
        }
        outputs = self.model(samples, blip2_samples, mask_infos=None, task_button=task_button)

        mask_infos = outputs['mask_infos_pred']
        pred_names = [list(mask_info.values()) for mask_info in mask_infos]
        orig_target_sizes = torch.tensor([tuple(reversed(image.size))]).to(device)
        results = self.postprocessor(outputs, orig_target_sizes, pred_names, mask_infos)[0]
        image_vis = visualize_prediction(image, results, threshold)

        out_text = outputs['output_text'][0]
        if task_button == "Cloze Test":
            history = []
            chat = [
                (prompt, out_text),
            ]
        elif task_button == "Captioning":
            history = []
            chat = [
                ("please describe the image", out_text),
            ]
        elif task_button == "Question Answering":
            history += [out_text]
            chat = [
                (history[i], history[i + 1]) for i in range(0, len(history) - 1, 2)
            ]
        else:
            history = []
            chat = []
        return image_vis, chat, history