aot/tools/demo.py

import importlib
import sys
import os

sys.path.append('.')
sys.path.append('..')

import cv2
from PIL import Image
from skimage.morphology.binary import binary_dilation

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms

from networks.models import build_vos_model
from networks.engines import build_engine
from utils.checkpoint import load_network

from dataloaders.eval_datasets import VOSTest
import dataloaders.video_transforms as tr
from utils.image import save_mask

_palette = [
    255, 0, 0, 0, 0, 139, 255, 255, 84, 0, 255, 0, 139, 0, 139, 0, 128, 128,
    128, 128, 128, 139, 0, 0, 218, 165, 32, 144, 238, 144, 160, 82, 45, 148, 0,
    211, 255, 0, 255, 30, 144, 255, 255, 218, 185, 85, 107, 47, 255, 140, 0,
    50, 205, 50, 123, 104, 238, 240, 230, 140, 72, 61, 139, 128, 128, 0, 0, 0,
    205, 221, 160, 221, 143, 188, 143, 127, 255, 212, 176, 224, 230, 244, 164,
    96, 250, 128, 114, 70, 130, 180, 0, 128, 0, 173, 255, 47, 255, 105, 180,
    238, 130, 238, 154, 205, 50, 220, 20, 60, 176, 48, 96, 0, 206, 209, 0, 191,
    255, 40, 40, 40, 41, 41, 41, 42, 42, 42, 43, 43, 43, 44, 44, 44, 45, 45,
    45, 46, 46, 46, 47, 47, 47, 48, 48, 48, 49, 49, 49, 50, 50, 50, 51, 51, 51,
    52, 52, 52, 53, 53, 53, 54, 54, 54, 55, 55, 55, 56, 56, 56, 57, 57, 57, 58,
    58, 58, 59, 59, 59, 60, 60, 60, 61, 61, 61, 62, 62, 62, 63, 63, 63, 64, 64,
    64, 65, 65, 65, 66, 66, 66, 67, 67, 67, 68, 68, 68, 69, 69, 69, 70, 70, 70,
    71, 71, 71, 72, 72, 72, 73, 73, 73, 74, 74, 74, 75, 75, 75, 76, 76, 76, 77,
    77, 77, 78, 78, 78, 79, 79, 79, 80, 80, 80, 81, 81, 81, 82, 82, 82, 83, 83,
    83, 84, 84, 84, 85, 85, 85, 86, 86, 86, 87, 87, 87, 88, 88, 88, 89, 89, 89,
    90, 90, 90, 91, 91, 91, 92, 92, 92, 93, 93, 93, 94, 94, 94, 95, 95, 95, 96,
    96, 96, 97, 97, 97, 98, 98, 98, 99, 99, 99, 100, 100, 100, 101, 101, 101,
    102, 102, 102, 103, 103, 103, 104, 104, 104, 105, 105, 105, 106, 106, 106,
    107, 107, 107, 108, 108, 108, 109, 109, 109, 110, 110, 110, 111, 111, 111,
    112, 112, 112, 113, 113, 113, 114, 114, 114, 115, 115, 115, 116, 116, 116,
    117, 117, 117, 118, 118, 118, 119, 119, 119, 120, 120, 120, 121, 121, 121,
    122, 122, 122, 123, 123, 123, 124, 124, 124, 125, 125, 125, 126, 126, 126,
    127, 127, 127, 128, 128, 128, 129, 129, 129, 130, 130, 130, 131, 131, 131,
    132, 132, 132, 133, 133, 133, 134, 134, 134, 135, 135, 135, 136, 136, 136,
    137, 137, 137, 138, 138, 138, 139, 139, 139, 140, 140, 140, 141, 141, 141,
    142, 142, 142, 143, 143, 143, 144, 144, 144, 145, 145, 145, 146, 146, 146,
    147, 147, 147, 148, 148, 148, 149, 149, 149, 150, 150, 150, 151, 151, 151,
    152, 152, 152, 153, 153, 153, 154, 154, 154, 155, 155, 155, 156, 156, 156,
    157, 157, 157, 158, 158, 158, 159, 159, 159, 160, 160, 160, 161, 161, 161,
    162, 162, 162, 163, 163, 163, 164, 164, 164, 165, 165, 165, 166, 166, 166,
    167, 167, 167, 168, 168, 168, 169, 169, 169, 170, 170, 170, 171, 171, 171,
    172, 172, 172, 173, 173, 173, 174, 174, 174, 175, 175, 175, 176, 176, 176,
    177, 177, 177, 178, 178, 178, 179, 179, 179, 180, 180, 180, 181, 181, 181,
    182, 182, 182, 183, 183, 183, 184, 184, 184, 185, 185, 185, 186, 186, 186,
    187, 187, 187, 188, 188, 188, 189, 189, 189, 190, 190, 190, 191, 191, 191,
    192, 192, 192, 193, 193, 193, 194, 194, 194, 195, 195, 195, 196, 196, 196,
    197, 197, 197, 198, 198, 198, 199, 199, 199, 200, 200, 200, 201, 201, 201,
    202, 202, 202, 203, 203, 203, 204, 204, 204, 205, 205, 205, 206, 206, 206,
    207, 207, 207, 208, 208, 208, 209, 209, 209, 210, 210, 210, 211, 211, 211,
    212, 212, 212, 213, 213, 213, 214, 214, 214, 215, 215, 215, 216, 216, 216,
    217, 217, 217, 218, 218, 218, 219, 219, 219, 220, 220, 220, 221, 221, 221,
    222, 222, 222, 223, 223, 223, 224, 224, 224, 225, 225, 225, 226, 226, 226,
    227, 227, 227, 228, 228, 228, 229, 229, 229, 230, 230, 230, 231, 231, 231,
    232, 232, 232, 233, 233, 233, 234, 234, 234, 235, 235, 235, 236, 236, 236,
    237, 237, 237, 238, 238, 238, 239, 239, 239, 240, 240, 240, 241, 241, 241,
    242, 242, 242, 243, 243, 243, 244, 244, 244, 245, 245, 245, 246, 246, 246,
    247, 247, 247, 248, 248, 248, 249, 249, 249, 250, 250, 250, 251, 251, 251,
    252, 252, 252, 253, 253, 253, 254, 254, 254, 255, 255, 255, 0, 0, 0
]
color_palette = np.array(_palette).reshape(-1, 3)


def overlay(image, mask, colors=[255, 0, 0], cscale=1, alpha=0.4):
    colors = np.atleast_2d(colors) * cscale

    im_overlay = image.copy()
    object_ids = np.unique(mask)

    for object_id in object_ids[1:]:
        # Overlay color on  binary mask

        foreground = image * alpha + np.ones(
            image.shape) * (1 - alpha) * np.array(colors[object_id])
        binary_mask = mask == object_id

        # Compose image
        im_overlay[binary_mask] = foreground[binary_mask]

        countours = binary_dilation(binary_mask) ^ binary_mask
        im_overlay[countours, :] = 0

    return im_overlay.astype(image.dtype)


def demo(cfg):
    video_fps = 15
    gpu_id = cfg.TEST_GPU_ID

    # Load pre-trained model
    print('Build AOT model.')
    model = build_vos_model(cfg.MODEL_VOS, cfg).cuda(gpu_id)

    print('Load checkpoint from {}'.format(cfg.TEST_CKPT_PATH))
    model, _ = load_network(model, cfg.TEST_CKPT_PATH, gpu_id)

    print('Build AOT engine.')
    engine = build_engine(cfg.MODEL_ENGINE,
                          phase='eval',
                          aot_model=model,
                          gpu_id=gpu_id,
                          long_term_mem_gap=cfg.TEST_LONG_TERM_MEM_GAP)

    # Prepare datasets for each sequence
    transform = transforms.Compose([
        tr.MultiRestrictSize(cfg.TEST_MIN_SIZE, cfg.TEST_MAX_SIZE,
                             cfg.TEST_FLIP, cfg.TEST_MULTISCALE,
                             cfg.MODEL_ALIGN_CORNERS),
        tr.MultiToTensor()
    ])
    image_root = os.path.join(cfg.TEST_DATA_PATH, 'images')
    label_root = os.path.join(cfg.TEST_DATA_PATH, 'masks')

    sequences = os.listdir(image_root)
    seq_datasets = []
    for seq_name in sequences:
        print('Build a dataset for sequence {}.'.format(seq_name))
        seq_images = np.sort(os.listdir(os.path.join(image_root, seq_name)))
        seq_labels = [seq_images[0].replace('jpg', 'png')]
        seq_dataset = VOSTest(image_root,
                              label_root,
                              seq_name,
                              seq_images,
                              seq_labels,
                              transform=transform)
        seq_datasets.append(seq_dataset)

    # Infer
    output_root = cfg.TEST_OUTPUT_PATH
    output_mask_root = os.path.join(output_root, 'pred_masks')
    if not os.path.exists(output_mask_root):
        os.makedirs(output_mask_root)

    for seq_dataset in seq_datasets:
        seq_name = seq_dataset.seq_name
        image_seq_root = os.path.join(image_root, seq_name)
        output_mask_seq_root = os.path.join(output_mask_root, seq_name)
        if not os.path.exists(output_mask_seq_root):
            os.makedirs(output_mask_seq_root)
        print('Build a dataloader for sequence {}.'.format(seq_name))
        seq_dataloader = DataLoader(seq_dataset,
                                    batch_size=1,
                                    shuffle=False,
                                    num_workers=cfg.TEST_WORKERS,
                                    pin_memory=True)

        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        output_video_path = os.path.join(
            output_root, '{}_{}fps.avi'.format(seq_name, video_fps))

        print('Start the inference of sequence {}:'.format(seq_name))
        model.eval()
        engine.restart_engine()
        with torch.no_grad():
            for frame_idx, samples in enumerate(seq_dataloader):
                sample = samples[0]
                img_name = sample['meta']['current_name'][0]

                obj_nums = sample['meta']['obj_num']
                output_height = sample['meta']['height']
                output_width = sample['meta']['width']
                obj_idx = sample['meta']['obj_idx']

                obj_nums = [int(obj_num) for obj_num in obj_nums]
                obj_idx = [int(_obj_idx) for _obj_idx in obj_idx]

                current_img = sample['current_img']
                current_img = current_img.cuda(gpu_id, non_blocking=True)

                if frame_idx == 0:
                    videoWriter = cv2.VideoWriter(
                        output_video_path, fourcc, video_fps,
                        (int(output_width), int(output_height)))
                    print(
                        'Object number: {}. Inference size: {}x{}. Output size: {}x{}.'
                        .format(obj_nums[0],
                                current_img.size()[2],
                                current_img.size()[3], int(output_height),
                                int(output_width)))
                    current_label = sample['current_label'].cuda(
                        gpu_id, non_blocking=True).float()
                    current_label = F.interpolate(current_label,
                                                  size=current_img.size()[2:],
                                                  mode="nearest")
                    # add reference frame
                    engine.add_reference_frame(current_img,
                                               current_label,
                                               frame_step=0,
                                               obj_nums=obj_nums)
                else:
                    print('Processing image {}...'.format(img_name))
                    # predict segmentation
                    engine.match_propogate_one_frame(current_img)
                    pred_logit = engine.decode_current_logits(
                        (output_height, output_width))
                    pred_prob = torch.softmax(pred_logit, dim=1)
                    pred_label = torch.argmax(pred_prob, dim=1,
                                              keepdim=True).float()
                    _pred_label = F.interpolate(pred_label,
                                                size=engine.input_size_2d,
                                                mode="nearest")
                    # update memory
                    engine.update_memory(_pred_label)

                    # save results
                    input_image_path = os.path.join(image_seq_root, img_name)
                    output_mask_path = os.path.join(
                        output_mask_seq_root,
                        img_name.split('.')[0] + '.png')

                    pred_label = Image.fromarray(
                        pred_label.squeeze(0).squeeze(0).cpu().numpy().astype(
                            'uint8')).convert('P')
                    pred_label.putpalette(_palette)
                    pred_label.save(output_mask_path)

                    input_image = Image.open(input_image_path)

                    overlayed_image = overlay(
                        np.array(input_image, dtype=np.uint8),
                        np.array(pred_label, dtype=np.uint8), color_palette)
                    videoWriter.write(overlayed_image[..., [2, 1, 0]])

        print('Save a visualization video to {}.'.format(output_video_path))
        videoWriter.release()


def main():
    import argparse
    parser = argparse.ArgumentParser(description="AOT Demo")
    parser.add_argument('--exp_name', type=str, default='default')

    parser.add_argument('--stage', type=str, default='pre_ytb_dav')
    parser.add_argument('--model', type=str, default='r50_aotl')

    parser.add_argument('--gpu_id', type=int, default=0)

    parser.add_argument('--data_path', type=str, default='./datasets/Demo')
    parser.add_argument('--output_path', type=str, default='./demo_output')
    parser.add_argument('--ckpt_path',
                        type=str,
                        default='./pretrain_models/R50_AOTL_PRE_YTB_DAV.pth')

    parser.add_argument('--max_resolution', type=float, default=480 * 1.3)

    parser.add_argument('--amp', action='store_true')
    parser.set_defaults(amp=False)

    args = parser.parse_args()

    engine_config = importlib.import_module('configs.' + args.stage)
    cfg = engine_config.EngineConfig(args.exp_name, args.model)

    cfg.TEST_GPU_ID = args.gpu_id

    cfg.TEST_CKPT_PATH = args.ckpt_path
    cfg.TEST_DATA_PATH = args.data_path
    cfg.TEST_OUTPUT_PATH = args.output_path

    cfg.TEST_MIN_SIZE = None
    cfg.TEST_MAX_SIZE = args.max_resolution * 800. / 480.

    if args.amp:
        with torch.cuda.amp.autocast(enabled=True):
            demo(cfg)
    else:
        demo(cfg)


if __name__ == '__main__':
    main()