|
|
import os
|
|
|
from os import path
|
|
|
|
|
|
from torch.utils.data.dataset import Dataset
|
|
|
from torchvision import transforms
|
|
|
from torchvision.transforms import InterpolationMode
|
|
|
import torch.nn.functional as F
|
|
|
from PIL import Image
|
|
|
import numpy as np
|
|
|
|
|
|
from dataset.range_transform import im_normalization
|
|
|
|
|
|
|
|
|
class VideoReader(Dataset):
|
|
|
"""
|
|
|
This class is used to read a video, one frame at a time
|
|
|
"""
|
|
|
def __init__(self, vid_name, image_dir, mask_dir, size=-1, to_save=None, use_all_mask=False, size_dir=None):
|
|
|
"""
|
|
|
image_dir - points to a directory of jpg images
|
|
|
mask_dir - points to a directory of png masks
|
|
|
size - resize min. side to size. Does nothing if <0.
|
|
|
to_save - optionally contains a list of file names without extensions
|
|
|
where the segmentation mask is required
|
|
|
use_all_mask - when true, read all available mask in mask_dir.
|
|
|
Default false. Set to true for YouTubeVOS validation.
|
|
|
"""
|
|
|
self.vid_name = vid_name
|
|
|
self.image_dir = image_dir
|
|
|
self.mask_dir = mask_dir
|
|
|
self.to_save = to_save
|
|
|
self.use_all_mask = use_all_mask
|
|
|
if size_dir is None:
|
|
|
self.size_dir = self.image_dir
|
|
|
else:
|
|
|
self.size_dir = size_dir
|
|
|
|
|
|
self.frames = sorted(os.listdir(self.image_dir))
|
|
|
self.palette = Image.open(path.join(mask_dir, sorted(os.listdir(mask_dir))[0])).getpalette()
|
|
|
self.first_gt_path = path.join(self.mask_dir, sorted(os.listdir(self.mask_dir))[0])
|
|
|
|
|
|
if size < 0:
|
|
|
self.im_transform = transforms.Compose([
|
|
|
transforms.ToTensor(),
|
|
|
im_normalization,
|
|
|
])
|
|
|
else:
|
|
|
self.im_transform = transforms.Compose([
|
|
|
transforms.ToTensor(),
|
|
|
im_normalization,
|
|
|
transforms.Resize(size, interpolation=InterpolationMode.BILINEAR),
|
|
|
])
|
|
|
self.size = size
|
|
|
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
|
frame = self.frames[idx]
|
|
|
info = {}
|
|
|
data = {}
|
|
|
info['frame'] = frame
|
|
|
info['save'] = (self.to_save is None) or (frame[:-4] in self.to_save)
|
|
|
|
|
|
im_path = path.join(self.image_dir, frame)
|
|
|
img = Image.open(im_path).convert('RGB')
|
|
|
|
|
|
if self.image_dir == self.size_dir:
|
|
|
shape = np.array(img).shape[:2]
|
|
|
else:
|
|
|
size_path = path.join(self.size_dir, frame)
|
|
|
size_im = Image.open(size_path).convert('RGB')
|
|
|
shape = np.array(size_im).shape[:2]
|
|
|
|
|
|
gt_path = path.join(self.mask_dir, frame[:-4]+'.png')
|
|
|
img = self.im_transform(img)
|
|
|
|
|
|
load_mask = self.use_all_mask or (gt_path == self.first_gt_path)
|
|
|
if load_mask and path.exists(gt_path):
|
|
|
mask = Image.open(gt_path).convert('P')
|
|
|
mask = np.array(mask, dtype=np.uint8)
|
|
|
data['mask'] = mask
|
|
|
|
|
|
info['shape'] = shape
|
|
|
info['need_resize'] = not (self.size < 0)
|
|
|
data['rgb'] = img
|
|
|
data['info'] = info
|
|
|
|
|
|
return data
|
|
|
|
|
|
def resize_mask(self, mask):
|
|
|
|
|
|
h, w = mask.shape[-2:]
|
|
|
min_hw = min(h, w)
|
|
|
return F.interpolate(mask, (int(h/min_hw*self.size), int(w/min_hw*self.size)),
|
|
|
mode='nearest')
|
|
|
|
|
|
def get_palette(self):
|
|
|
return self.palette
|
|
|
|
|
|
def __len__(self):
|
|
|
return len(self.frames) |
