| import torch |
| import os |
| import json |
| import numpy as np |
| import random |
| from torch.utils.data import Dataset |
| from PIL import Image |
| from torchvision import transforms |
| from torchvision.transforms import functional as F |
|
|
|
|
| class YouTubeVISDataset(Dataset): |
| def __init__(self, videodir, annfile, size, seq_length, seq_sampler, transform=None): |
| self.videodir = videodir |
| self.size = size |
| self.seq_length = seq_length |
| self.seq_sampler = seq_sampler |
| self.transform = transform |
| |
| with open(annfile) as f: |
| data = json.load(f) |
|
|
| self.masks = {} |
| for ann in data['annotations']: |
| if ann['category_id'] == 26: |
| video_id = ann['video_id'] |
| if video_id not in self.masks: |
| self.masks[video_id] = [[] for _ in range(len(ann['segmentations']))] |
| for frame, mask in zip(self.masks[video_id], ann['segmentations']): |
| if mask is not None: |
| frame.append(mask) |
| |
| self.videos = {} |
| for video in data['videos']: |
| video_id = video['id'] |
| if video_id in self.masks: |
| self.videos[video_id] = video |
| |
| self.index = [] |
| for video_id in self.videos.keys(): |
| for frame in range(len(self.videos[video_id]['file_names'])): |
| self.index.append((video_id, frame)) |
| |
| def __len__(self): |
| return len(self.index) |
| |
| def __getitem__(self, idx): |
| video_id, frame_id = self.index[idx] |
| video = self.videos[video_id] |
| frame_count = len(self.videos[video_id]['file_names']) |
| H, W = video['height'], video['width'] |
| |
| imgs, segs = [], [] |
| for t in self.seq_sampler(self.seq_length): |
| frame = (frame_id + t) % frame_count |
|
|
| filename = video['file_names'][frame] |
| masks = self.masks[video_id][frame] |
| |
| with Image.open(os.path.join(self.videodir, filename)) as img: |
| imgs.append(self._downsample_if_needed(img.convert('RGB'), Image.BILINEAR)) |
| |
| seg = np.zeros((H, W), dtype=np.uint8) |
| for mask in masks: |
| seg |= self._decode_rle(mask) |
| segs.append(self._downsample_if_needed(Image.fromarray(seg), Image.NEAREST)) |
| |
| if self.transform is not None: |
| imgs, segs = self.transform(imgs, segs) |
| |
| return imgs, segs |
| |
| def _decode_rle(self, rle): |
| H, W = rle['size'] |
| msk = np.zeros(H * W, dtype=np.uint8) |
| encoding = rle['counts'] |
| skip = 0 |
| for i in range(0, len(encoding) - 1, 2): |
| skip += encoding[i] |
| draw = encoding[i + 1] |
| msk[skip : skip + draw] = 255 |
| skip += draw |
| return msk.reshape(W, H).transpose() |
| |
| def _downsample_if_needed(self, img, resample): |
| w, h = img.size |
| if min(w, h) > self.size: |
| scale = self.size / min(w, h) |
| w = int(scale * w) |
| h = int(scale * h) |
| img = img.resize((w, h), resample) |
| return img |
|
|
|
|
| class YouTubeVISAugmentation: |
| def __init__(self, size): |
| self.size = size |
| self.jitter = transforms.ColorJitter(0.3, 0.3, 0.3, 0.15) |
| |
| def __call__(self, imgs, segs): |
| |
| |
| imgs = torch.stack([F.to_tensor(img) for img in imgs]) |
| segs = torch.stack([F.to_tensor(seg) for seg in segs]) |
| |
| |
| params = transforms.RandomResizedCrop.get_params(imgs, scale=(0.8, 1), ratio=(0.9, 1.1)) |
| imgs = F.resized_crop(imgs, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) |
| segs = F.resized_crop(segs, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) |
| |
| |
| imgs = self.jitter(imgs) |
| |
| |
| if random.random() < 0.05: |
| imgs = F.rgb_to_grayscale(imgs, num_output_channels=3) |
| |
| |
| if random.random() < 0.5: |
| imgs = F.hflip(imgs) |
| segs = F.hflip(segs) |
| |
| return imgs, segs |
|
|
