3DFauna_demo / video3d /dataloaders.py
kyleleey
first commit
98a77e0
import os
from glob import glob
import random
import numpy as np
from PIL import Image
import cv2
import torch
from torch.utils.data import Dataset
import torchvision.datasets.folder
import torchvision.transforms as transforms
from einops import rearrange
def compute_distance_transform(mask):
mask_dt = []
for m in mask:
dt = torch.FloatTensor(cv2.distanceTransform(np.uint8(m[0]), cv2.DIST_L2, cv2.DIST_MASK_PRECISE))
inv_dt = torch.FloatTensor(cv2.distanceTransform(np.uint8(1 - m[0]), cv2.DIST_L2, cv2.DIST_MASK_PRECISE))
mask_dt += [torch.stack([dt, inv_dt], 0)]
return torch.stack(mask_dt, 0) # Bx2xHxW
def crop_image(image, boxs, size):
crops = []
for box in boxs:
crop_x0, crop_y0, crop_w, crop_h = box
crop = transforms.functional.resized_crop(image, crop_y0, crop_x0, crop_h, crop_w, size)
crop = transforms.functional.to_tensor(crop)
crops += [crop]
return torch.stack(crops, 0)
def box_loader(fpath):
box = np.loadtxt(fpath, 'str')
box[0] = box[0].split('_')[0]
return box.astype(np.float32)
def read_feat_from_img(path, n_channels):
feat = np.array(Image.open(path))
return dencode_feat_from_img(feat, n_channels)
def dencode_feat_from_img(img, n_channels):
n_addon_channels = int(np.ceil(n_channels / 3) * 3) - n_channels
n_tiles = int((n_channels + n_addon_channels) / 3)
feat = rearrange(img, 'h (t w) c -> h w (t c)', t=n_tiles, c=3)
feat = feat[:, :, :-n_addon_channels]
feat = feat.astype('float32') / 255
return feat.transpose(2, 0, 1)
def dino_loader(fpath, n_channels):
dino_map = read_feat_from_img(fpath, n_channels)
return dino_map
def get_valid_mask(boxs, image_size):
valid_masks = []
for box in boxs:
crop_x0, crop_y0, crop_w, crop_h, full_w, full_h = box[1:7].int().numpy()
# Discard a small margin near the boundary.
margin_w = int(crop_w * 0.02)
margin_h = int(crop_h * 0.02)
mask_full = torch.ones(full_h-margin_h*2, full_w-margin_w*2)
mask_full_pad = torch.nn.functional.pad(mask_full, (crop_w+margin_w, crop_w+margin_w, crop_h+margin_h, crop_h+margin_h), mode='constant', value=0.0)
mask_full_crop = mask_full_pad[crop_y0+crop_h:crop_y0+crop_h*2, crop_x0+crop_w:crop_x0+crop_w*2]
mask_crop = torch.nn.functional.interpolate(mask_full_crop[None, None, :, :], image_size, mode='nearest')[0,0]
valid_masks += [mask_crop]
return torch.stack(valid_masks, 0) # NxHxW
def horizontal_flip_box(box):
frame_id, crop_x0, crop_y0, crop_w, crop_h, full_w, full_h, sharpness, label = box.unbind(1)
box[:,1] = full_w - crop_x0 - crop_w # x0
return box
def horizontal_flip_all(images, masks, mask_dt, mask_valid, flows, bboxs, bg_images, dino_features=None, dino_clusters=None):
images = images.flip(3) # NxCxHxW
masks = masks.flip(3) # NxCxHxW
mask_dt = mask_dt.flip(3) # NxCxHxW
mask_valid = mask_valid.flip(2) # NxHxW
if flows.dim() > 1:
flows = flows.flip(3) # (N-1)x(x,y)xHxW
flows[:,0] *= -1 # invert delta x
bboxs = horizontal_flip_box(bboxs) # NxK
bg_images = bg_images.flip(3) # NxCxHxW
if dino_features.dim() > 1:
dino_features = dino_features.flip(3)
if dino_clusters.dim() > 1:
dino_clusters = dino_clusters.flip(3)
return images, masks, mask_dt, mask_valid, flows, bboxs, bg_images, dino_features, dino_clusters
class BaseSequenceDataset(Dataset):
def __init__(self, root, skip_beginning=4, skip_end=4, min_seq_len=10, debug_seq=False):
super().__init__()
self.skip_beginning = skip_beginning
self.skip_end = skip_end
self.min_seq_len = min_seq_len
# self.pattern = "{:07d}_{}"
self.sequences = self._make_sequences(root)
if debug_seq:
# self.sequences = [self.sequences[0][20:160]] * 100
seq_len = 0
while seq_len < min_seq_len:
i = np.random.randint(len(self.sequences))
rand_seq = self.sequences[i]
seq_len = len(rand_seq)
self.sequences = [rand_seq]
self.samples = []
def _make_sequences(self, path):
result = []
for d in sorted(os.scandir(path), key=lambda e: e.name):
if d.is_dir():
files = self._parse_folder(d)
if len(files) >= self.min_seq_len:
result.append(files)
return result
def _parse_folder(self, path):
result = sorted(glob(os.path.join(path, '*'+self.image_loaders[0][0])))
result = [p.replace(self.image_loaders[0][0], '{}') for p in result]
if len(result) <= self.skip_beginning + self.skip_end:
return []
if self.skip_end == 0:
return result[self.skip_beginning:]
return result[self.skip_beginning:-self.skip_end]
def _load_ids(self, path_patterns, loaders, transform=None):
result = []
for loader in loaders:
for p in path_patterns:
x = loader[1](p.format(loader[0]), *loader[2:])
if transform:
x = transform(x)
result.append(x)
return tuple(result)
def __len__(self):
return len(self.samples)
def __getitem__(self, index):
raise NotImplemented("This is a base class and should not be used directly")
class NFrameSequenceDataset(BaseSequenceDataset):
def __init__(self, root, cat_name=None, num_sample_frames=2, skip_beginning=4, skip_end=4, min_seq_len=10, in_image_size=256, out_image_size=256, debug_seq=False, random_sample=False, shuffle=False, dense_sample=True, color_jitter=None, load_background=False, random_flip=False, rgb_suffix='.png', load_dino_feature=False, load_dino_cluster=False, dino_feature_dim=64, **kwargs):
self.cat_name = cat_name
self.image_loaders = [("rgb"+rgb_suffix, torchvision.datasets.folder.default_loader)]
self.mask_loaders = [("mask.png", torchvision.datasets.folder.default_loader)]
self.bbox_loaders = [("box.txt", box_loader)]
super().__init__(root, skip_beginning, skip_end, min_seq_len, debug_seq)
if num_sample_frames > 1:
self.flow_loaders = [("flow.png", cv2.imread, cv2.IMREAD_UNCHANGED)]
else:
self.flow_loaders = None
self.num_sample_frames = num_sample_frames
self.random_sample = random_sample
if self.random_sample:
if shuffle:
random.shuffle(self.sequences)
self.samples = self.sequences
else:
for i, s in enumerate(self.sequences):
stride = 1 if dense_sample else self.num_sample_frames
self.samples += [(i, k) for k in range(0, len(s), stride)]
if shuffle:
random.shuffle(self.samples)
self.in_image_size = in_image_size
self.out_image_size = out_image_size
self.load_background = load_background
self.color_jitter = color_jitter
self.image_transform = transforms.Compose([transforms.Resize(self.in_image_size), transforms.ToTensor()])
self.mask_transform = transforms.Compose([transforms.Resize(self.out_image_size, interpolation=Image.NEAREST), transforms.ToTensor()])
if self.flow_loaders is not None:
self.flow_transform = lambda x: (torch.FloatTensor(x.astype(np.float32)).flip(2)[:,:,:2] / 65535. ) *2 -1
self.random_flip = random_flip
self.load_dino_feature = load_dino_feature
if load_dino_feature:
self.dino_feature_loaders = [(f"feat{dino_feature_dim}.png", dino_loader, dino_feature_dim)]
self.load_dino_cluster = load_dino_cluster
if load_dino_cluster:
self.dino_cluster_loaders = [("clusters.png", torchvision.datasets.folder.default_loader)]
def __getitem__(self, index):
if self.random_sample:
seq_idx = index % len(self.sequences)
seq = self.sequences[seq_idx]
if len(seq) < self.num_sample_frames:
start_frame_idx = 0
else:
start_frame_idx = np.random.randint(len(seq)-self.num_sample_frames+1)
paths = seq[start_frame_idx:start_frame_idx+self.num_sample_frames]
else:
seq_idx, start_frame_idx = self.samples[index % len(self.samples)]
seq = self.sequences[seq_idx]
# Handle edge case: when only last frame is left, sample last two frames, except if the sequence only has one frame
if len(seq) <= start_frame_idx +1:
start_frame_idx = max(0, start_frame_idx-1)
paths = seq[start_frame_idx:start_frame_idx+self.num_sample_frames]
masks = torch.stack(self._load_ids(paths, self.mask_loaders, transform=self.mask_transform), 0) # load all images
mask_dt = compute_distance_transform(masks)
jitter = False
if self.color_jitter is not None:
prob, b, h = self.color_jitter
if np.random.rand() < prob:
jitter = True
color_jitter_tsf_fg = transforms.ColorJitter.get_params(brightness=(1-b, 1+b), contrast=None, saturation=None, hue=(-h, h))
image_transform_fg = transforms.Compose([transforms.Resize(self.in_image_size), color_jitter_tsf_fg, transforms.ToTensor()])
color_jitter_tsf_bg = transforms.ColorJitter.get_params(brightness=(1-b, 1+b), contrast=None, saturation=None, hue=(-h, h))
image_transform_bg = transforms.Compose([transforms.Resize(self.in_image_size), color_jitter_tsf_bg, transforms.ToTensor()])
if jitter:
images_fg = torch.stack(self._load_ids(paths, self.image_loaders, transform=image_transform_fg), 0) # load all images
images_bg = torch.stack(self._load_ids(paths, self.image_loaders, transform=image_transform_bg), 0) # load all images
images = images_fg * masks + images_bg * (1-masks)
else:
images = torch.stack(self._load_ids(paths, self.image_loaders, transform=self.image_transform), 0) # load all images
if len(paths) > 1:
flows = torch.stack(self._load_ids(paths[:-1], self.flow_loaders, transform=self.flow_transform), 0).permute(0,3,1,2) # load flow for first image, (N-1)x(x,y)xHxW, -1~1
flows = torch.nn.functional.interpolate(flows, size=self.out_image_size, mode="bilinear")
else:
flows = torch.zeros(1)
bboxs = torch.stack(self._load_ids(paths, self.bbox_loaders, transform=torch.FloatTensor), 0) # load bounding boxes for all images
mask_valid = get_valid_mask(bboxs, (self.out_image_size, self.out_image_size)) # exclude pixels cropped outside the original image
if self.load_background:
bg_image = torchvision.datasets.folder.default_loader(os.path.join(os.path.dirname(paths[0]), 'background_frame.jpg'))
if jitter:
bg_image = color_jitter_tsf_bg(bg_image)
bg_images = crop_image(bg_image, bboxs[:, 1:5].int().numpy(), (self.in_image_size, self.in_image_size))
else:
bg_images = torch.zeros_like(images)
if self.load_dino_feature:
dino_features = torch.stack(self._load_ids(paths, self.dino_feature_loaders, transform=torch.FloatTensor), 0) # BxFx64x224x224
else:
dino_features = torch.zeros(1)
if self.load_dino_cluster:
dino_clusters = torch.stack(self._load_ids(paths, self.dino_cluster_loaders, transform=transforms.ToTensor()), 0) # BxFx3x55x55
else:
dino_clusters = torch.zeros(1)
seq_idx = torch.LongTensor([seq_idx])
frame_idx = torch.arange(start_frame_idx, start_frame_idx+len(paths)).long()
if self.random_flip and np.random.rand() < 0.5:
images, masks, mask_dt, mask_valid, flows, bboxs, bg_images, dino_features, dino_clusters = horizontal_flip_all(images, masks, mask_dt, mask_valid, flows, bboxs, bg_images, dino_features, dino_clusters)
## pad shorter sequence
if len(paths) < self.num_sample_frames:
num_pad = self.num_sample_frames - len(paths)
images = torch.cat([images[:1]] *num_pad + [images], 0)
masks = torch.cat([masks[:1]] *num_pad + [masks], 0)
mask_dt = torch.cat([mask_dt[:1]] *num_pad + [mask_dt], 0)
mask_valid = torch.cat([mask_valid[:1]] *num_pad + [mask_valid], 0)
if flows.dim() > 1:
flows = torch.cat([flows[:1]*0] *num_pad + [flows], 0)
bboxs = torch.cat([bboxs[:1]] * num_pad + [bboxs], 0)
bg_images = torch.cat([bg_images[:1]] *num_pad + [bg_images], 0)
if dino_features.dim() > 1:
dino_features = torch.cat([dino_features[:1]] *num_pad + [dino_features], 0)
if dino_clusters.dim() > 1:
dino_clusters = torch.cat([dino_clusters[:1]] *num_pad + [dino_clusters], 0)
frame_idx = torch.cat([frame_idx[:1]] *num_pad + [frame_idx], 0)
return images, masks, mask_dt, mask_valid, flows, bboxs, bg_images, dino_features, dino_clusters, seq_idx, frame_idx, self.cat_name
def get_sequence_loader(data_dir, **kwargs):
if isinstance(data_dir, dict):
loaders = []
for k, v in data_dir.items():
dataset= NFrameSequenceDataset(v, cat_name=k, **kwargs)
loader = torch.utils.data.DataLoader(dataset, batch_size=kwargs['batch_size'], shuffle=kwargs['shuffle'], num_workers=kwargs['num_workers'], pin_memory=True)
loaders += [loader]
return loaders
else:
return [get_sequence_loader_single(data_dir, **kwargs)]
def get_sequence_loader_single(data_dir, mode='all_frame', is_validation=False, batch_size=256, num_workers=4, in_image_size=256, out_image_size=256, debug_seq=False, num_sample_frames=2, skip_beginning=4, skip_end=4, min_seq_len=10, max_seq_len=256, random_sample=False, shuffle=False, dense_sample=True, color_jitter=None, load_background=False, random_flip=False, rgb_suffix='.jpg', load_dino_feature=False, load_dino_cluster=False, dino_feature_dim=64):
if mode == 'n_frame':
dataset = NFrameSequenceDataset(data_dir, num_sample_frames=num_sample_frames, skip_beginning=skip_beginning, skip_end=skip_end, min_seq_len=min_seq_len, in_image_size=in_image_size, out_image_size=out_image_size, debug_seq=debug_seq, random_sample=random_sample, shuffle=shuffle, dense_sample=dense_sample, color_jitter=color_jitter, load_background=load_background, random_flip=random_flip, rgb_suffix=rgb_suffix, load_dino_feature=load_dino_feature, load_dino_cluster=load_dino_cluster, dino_feature_dim=dino_feature_dim)
else:
raise NotImplementedError
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
shuffle=not is_validation,
num_workers=num_workers,
pin_memory=True
)
return loader
class ImageDataset(Dataset):
def __init__(self, root, is_validation=False, image_size=256, color_jitter=None):
super().__init__()
self.image_loader = ("rgb.jpg", torchvision.datasets.folder.default_loader)
self.mask_loader = ("mask.png", torchvision.datasets.folder.default_loader)
self.bbox_loader = ("box.txt", np.loadtxt, 'str')
self.samples = self._parse_folder(root)
self.image_size = image_size
self.color_jitter = color_jitter
self.image_transform = transforms.Compose([transforms.Resize(self.image_size), transforms.ToTensor()])
self.mask_transform = transforms.Compose([transforms.Resize(self.image_size, interpolation=Image.NEAREST), transforms.ToTensor()])
def _parse_folder(self, path):
result = sorted(glob(os.path.join(path, '**/*'+self.image_loader[0]), recursive=True))
result = [p.replace(self.image_loader[0], '{}') for p in result]
return result
def _load_ids(self, path, loader, transform=None):
x = loader[1](path.format(loader[0]), *loader[2:])
if transform:
x = transform(x)
return x
def __len__(self):
return len(self.samples)
def __getitem__(self, index):
path = self.samples[index % len(self.samples)]
masks = self._load_ids(path, self.mask_loader, transform=self.mask_transform).unsqueeze(0)
mask_dt = compute_distance_transform(masks)
jitter = False
if self.color_jitter is not None:
prob, b, h = self.color_jitter
if np.random.rand() < prob:
jitter = True
color_jitter_tsf_fg = transforms.ColorJitter.get_params(brightness=(1-b, 1+b), contrast=None, saturation=None, hue=(-h, h))
image_transform_fg = transforms.Compose([transforms.Resize(self.image_size), color_jitter_tsf_fg, transforms.ToTensor()])
color_jitter_tsf_bg = transforms.ColorJitter.get_params(brightness=(1-b, 1+b), contrast=None, saturation=None, hue=(-h, h))
image_transform_bg = transforms.Compose([transforms.Resize(self.image_size), color_jitter_tsf_bg, transforms.ToTensor()])
if jitter:
images_fg = self._load_ids(path, self.image_loader, transform=image_transform_fg).unsqueeze(0)
images_bg = self._load_ids(path, self.image_loader, transform=image_transform_bg).unsqueeze(0)
images = images_fg * masks + images_bg * (1-masks)
else:
images = self._load_ids(path, self.image_loader, transform=self.image_transform).unsqueeze(0)
flows = torch.zeros(1)
bboxs = self._load_ids(path, self.bbox_loader, transform=None)
bboxs[0] = '0'
bboxs = torch.FloatTensor(bboxs.astype('float')).unsqueeze(0)
bg_fpath = os.path.join(os.path.dirname(path), 'background_frame.jpg')
if os.path.isfile(bg_fpath):
bg_image = torchvision.datasets.folder.default_loader(bg_fpath)
if jitter:
bg_image = color_jitter_tsf_bg(bg_image)
bg_image = transforms.ToTensor()(bg_image)
else:
bg_image = images[0]
seq_idx = torch.LongTensor([index])
frame_idx = torch.LongTensor([0])
return images, masks, mask_dt, flows, bboxs, bg_image, seq_idx, frame_idx
def get_image_loader(data_dir, is_validation=False, batch_size=256, num_workers=4, image_size=256, color_jitter=None):
dataset = ImageDataset(data_dir, is_validation=is_validation, image_size=image_size, color_jitter=color_jitter)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=True
)
return loader