eval_youtube_phase2.py

"""
YouTubeVOS has a label structure that is more complicated than DAVIS
Labels might not appear on the first frame (there might be no labels at all in the first frame)
Labels might not even appear on the same frame (i.e. Object 0 at frame 10, and object 1 at frame 15)
0 does not mean background -- it is simply "no-label"
and object indices might not be in order, there are missing indices somewhere in the validation set

Dealing with these makes the logic a bit convoluted here
It is not necessarily hacky but do understand that it is not as straightforward as DAVIS

Validation/test set.
"""


import os
from os import path
from argparse import ArgumentParser
import json

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np
from PIL import Image

from model.eval_network import STCN
from dataset.yv_test_dataset import YouTubeVOSTestDatasetPhase2
from util.tensor_util import unpad
from inference_core_yv import InferenceCore

from progressbar import progressbar
import sys

"""
Arguments loading
"""
parser = ArgumentParser()
parser.add_argument('--model', default='saves/stcn.pth')
parser.add_argument('--yv_path', default='../YouTube')
parser.add_argument('--davis_path', default = None)

parser.add_argument('--output_all', help=
"""
We will output all the frames if this is set to true.
Otherwise only a subset will be outputted, as determined by meta.json to save disk space.
For ensemble, all the sources must have this setting unified.
""", action='store_true')

parser.add_argument('--output')
parser.add_argument('--data_file', help='valid/test', default='valid')
parser.add_argument('--top', type=int, default=20)
parser.add_argument('--amp', action='store_true')
parser.add_argument('--mem_every', default=1, type=int)
parser.add_argument('--include_last', help='include last frame as temporary memory?', action='store_true')
parser.add_argument('--reverse', action = 'store_true')
parser.add_argument('--only_sec_gt', action = 'store_true')
parser.add_argument('--davis', action = 'store_true')
parser.add_argument('--res', default = 480, type = int )
args = parser.parse_args()

yv_path = args.yv_path
out_path = args.output


# Simple setup
os.makedirs(out_path, exist_ok=True)
if not args.davis:
    print(f'res is {args.res}')
    if args.res == -1:
        palette = Image.open(path.expanduser(yv_path + '/train_480p/Annotations/fffe5f8df6/00000.png')).getpalette()
    else:
        palette = Image.open(path.expanduser(yv_path + '/train/Annotations/fffe5f8df6/00000.png')).getpalette()
else:
    if args.davis_path == None:
        args.davis_path = '../DAVIS/2017'
        davis_path = args.davis_path
        print('davis_path is set to ../DAVIS/2017')
    palette = Image.open(path.expanduser(davis_path + '/trainval/Annotations/480p/blackswan/00000.png')).getpalette()

torch.autograd.set_grad_enabled(False)

# Load the json if we have to
if not args.output_all:
    with open(path.join(yv_path, args.split, 'meta.json')) as f:
        meta = json.load(f)['videos']

# Setup Dataset
if not args.davis:
    test_dataset = YouTubeVOSTestDatasetPhase2(data_root=yv_path, data_file=args.data_file,
                    reverse = args.reverse, only_sec_gt = args.only_sec_gt, res = args.res)
else:
    test_dataset = YouTubeVOSTestDatasetPhase2(data_root=davis_path, data_file=args.data_file,
                    reverse = args.reverse, only_sec_gt = args.only_sec_gt, res = -1, davis = True)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=4)

# Load our checkpoint
top_k = args.top
prop_model = STCN().cuda().eval()

# Performs input mapping such that stage 0 model can be loaded
print(f'model: {args.model}')
prop_saved = torch.load(args.model)
for k in list(prop_saved.keys()):
    if k == 'value_encoder.conv1.weight':
        if prop_saved[k].shape[1] == 4:
            pads = torch.zeros((64,1,7,7), device=prop_saved[k].device)
            prop_saved[k] = torch.cat([prop_saved[k], pads], 1)
prop_model.load_state_dict(prop_saved)

# Start eval
for data in progressbar(test_loader, max_value=len(test_loader), redirect_stdout=True):
    with torch.cuda.amp.autocast(enabled=args.amp):
        rgb = data['rgb']
        msk = data['gt'][0]
        info = data['info']
        name = info['name'][0]
        num_objects = len(info['labels'][0])
        gt_obj = info['gt_obj']
        size = info['size']

        # Load the required set of frames (if we don't need all)
        req_frames = None
        if not args.output_all:
            req_frames = []
            objects = meta[name]['objects']
            for key, value in objects.items():
                req_frames.extend(value['frames'])

            # Map the frame names to indices
            req_frames_names = set(req_frames)
            req_frames = []
            for fi in range(rgb.shape[1]):
                frame_name = info['frames'][fi][0][:-4]
                if frame_name in req_frames_names:
                    req_frames.append(fi)
            req_frames = sorted(req_frames)

        # Frames with labels, but they are not exhaustively labeled
        frames_with_gt = sorted(list(gt_obj.keys())) # [0]

        processor = InferenceCore(prop_model, rgb, num_objects=num_objects, top_k=top_k,
                                    mem_every=args.mem_every, include_last=args.include_last,
                                    req_frames=req_frames) # top_k: 20, mem_every: 5

        # min_idx tells us the starting point of propagation
        # Propagating before there are labels is not useful
        min_idx = 99999
        for i, frame_idx in enumerate(frames_with_gt):
            min_idx = min(frame_idx, min_idx)
            # Note that there might be more than one label per frame
            obj_idx = gt_obj[frame_idx][0].tolist()

            # Map the possibly non-continuous labels into a continuous scheme
            obj_idx = [info['label_convert'][o].item() for o in obj_idx]

            # Append the background label
            with_bg_msk = torch.cat([
                1 - torch.sum(msk[:,frame_idx], dim=0, keepdim=True),
                msk[:,frame_idx],
            ], 0).cuda()

            # We perform propagation from the current frame to the next frame with label
            if i == len(frames_with_gt) - 1:
                processor.interact(with_bg_msk, frame_idx, rgb.shape[1], obj_idx)
            else:
                processor.interact(with_bg_msk, frame_idx, frames_with_gt[i+1]+1, obj_idx)

        # Do unpad -> upsample to original size (we made it 480p)
        out_masks = torch.zeros((processor.t, 1, *size), dtype=torch.uint8, device='cuda')

        for ti in range(processor.t):
            prob = unpad(processor.prob[:,ti], processor.pad)
            prob = F.interpolate(prob, size, mode='bilinear', align_corners=False)
            out_masks[ti] = torch.argmax(prob, dim=0)

        out_masks = (out_masks.detach().cpu().numpy()[:,0]).astype(np.uint8)

        # Remap the indices to the original domain
        idx_masks = np.zeros_like(out_masks)
        for i in range(1, num_objects+1):
            backward_idx = info['label_backward'][i].item()
            idx_masks[out_masks==i] = backward_idx

        # Save the results
        this_out_path = path.join(out_path, 'Annotations', name)
        os.makedirs(this_out_path, exist_ok=True)
        for f in range(idx_masks.shape[0]):
            if f >= min_idx:
                if args.output_all or (f in req_frames):
                    img_E = Image.fromarray(idx_masks[f])
                    img_E.putpalette(palette)
                    img_E.save(os.path.join(this_out_path, info['frames'][f][0].replace('.jpg','.png')))

        del rgb
        del msk
        del processor