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()