713 lines
24 KiB
Python
713 lines
24 KiB
Python
import json
|
|
import os
|
|
import random
|
|
from dataclasses import dataclass, field
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
from ..utils.config import parse_structured
|
|
from ..utils.typing import *
|
|
|
|
|
|
def _parse_object_list_single(object_list_path: str):
|
|
all_objects = []
|
|
if object_list_path.endswith(".json"):
|
|
with open(object_list_path) as f:
|
|
all_objects = json.loads(f.read())
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
return all_objects
|
|
|
|
|
|
def _parse_object_list(object_list_path: Union[str, List[str]]):
|
|
all_objects = []
|
|
if isinstance(object_list_path, str):
|
|
object_list_path = [object_list_path]
|
|
for object_list_path_ in object_list_path:
|
|
all_objects += _parse_object_list_single(object_list_path_)
|
|
return all_objects
|
|
|
|
|
|
def _parse_scene_list_single(scene_list_path: str, root_data_dir: str):
|
|
all_scenes = []
|
|
if scene_list_path.endswith(".json"):
|
|
with open(scene_list_path) as f:
|
|
for p in json.loads(f.read()):
|
|
all_scenes.append(os.path.join(root_data_dir, p))
|
|
elif scene_list_path.endswith(".txt"):
|
|
with open(scene_list_path) as f:
|
|
for p in f.readlines():
|
|
p = p.strip()
|
|
all_scenes.append(os.path.join(root_data_dir, p))
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
return all_scenes
|
|
|
|
|
|
def _parse_scene_list(
|
|
scene_list_path: Union[str, List[str]], root_data_dir: Union[str, List[str]]
|
|
):
|
|
all_scenes = []
|
|
if isinstance(scene_list_path, str):
|
|
scene_list_path = [scene_list_path]
|
|
if isinstance(root_data_dir, str):
|
|
root_data_dir = [root_data_dir]
|
|
for scene_list_path_, root_data_dir_ in zip(scene_list_path, root_data_dir):
|
|
all_scenes += _parse_scene_list_single(scene_list_path_, root_data_dir_)
|
|
return all_scenes
|
|
|
|
|
|
def random_morphological_transform(
|
|
mask: torch.Tensor, max_kernel_size: int = 5, p_dilation: float = 0.5
|
|
) -> torch.Tensor:
|
|
"""
|
|
Randomly dilate or erode the mask
|
|
:param mask: [H, W] 0-1 mask
|
|
:param max_kernel_size: Maximum kernel size; controls the scale of the morphological operation
|
|
:param p_dilation: Probability of dilation; if random sample > p_dilation, erosion is performed
|
|
:return: Perturbed mask
|
|
"""
|
|
mask_np = mask.cpu().numpy().astype(np.uint8)
|
|
|
|
kernel_size = np.random.randint(1, max_kernel_size + 1)
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
kernel = np.ones((kernel_size, kernel_size), np.uint8)
|
|
|
|
if np.random.rand() < p_dilation:
|
|
mask_np = cv2.dilate(mask_np, kernel, iterations=1)
|
|
else:
|
|
mask_np = cv2.erode(mask_np, kernel, iterations=1)
|
|
|
|
return torch.from_numpy(mask_np).float().to(mask.device)
|
|
|
|
|
|
@dataclass
|
|
class MultiObjectDataModuleConfig:
|
|
scene_list: Any = ""
|
|
object_list: Any = ""
|
|
|
|
# Surface
|
|
surface_root_dir: Any = ""
|
|
surface_suffix: str = "npy"
|
|
num_surface_samples_per_object: int = 20480
|
|
return_scene: bool = False
|
|
|
|
max_num_instances: Optional[int] = None
|
|
padding: bool = True
|
|
num_instances_per_batch: Optional[int] = 10
|
|
|
|
# Image input
|
|
image_root_dir: Any = ""
|
|
image_prefix: Any = "render"
|
|
image_suffix: str = "webp"
|
|
idmap_prefix: str = "semantic"
|
|
idmap_suffix: str = "png"
|
|
background_color: Union[str, float] = "white"
|
|
image_names: List[str] = field(default_factory=lambda: [])
|
|
height: int = 768
|
|
width: int = 768
|
|
|
|
use_scene_image: bool = True
|
|
remove_scene_bg: bool = False
|
|
|
|
# Data processing
|
|
skip_small_object: bool = False
|
|
small_image_proportion: float = 0.005 # (16/224)^2
|
|
|
|
## Mask perturbation
|
|
morph_perturb: bool = False
|
|
max_kernel_size: int = 5
|
|
p_dilation: float = 0.5
|
|
|
|
return_crop_padded: bool = False
|
|
height_crop_padded: int = 224
|
|
width_crop_padded: int = 224
|
|
|
|
# Mix data
|
|
do_mix: bool = False
|
|
do_mix_prob: float = 0.5
|
|
|
|
mix_length: int = 80000
|
|
mix_scene_list: str = ""
|
|
mix_image_root_dir: str = ""
|
|
mix_surface_root_dir: str = ""
|
|
mix_surface_suffix: str = "npy"
|
|
mix_image_prefix: str = "render_opaque"
|
|
mix_image_names: List[str] = field(default_factory=lambda: [])
|
|
mix_image_suffix: str = "webp"
|
|
|
|
train_indices: Optional[Tuple[Any, Any]] = None
|
|
val_indices: Optional[Tuple[Any, Any]] = None
|
|
test_indices: Optional[Tuple[Any, Any]] = None
|
|
|
|
repeat: int = 1
|
|
|
|
batch_size: int = 1
|
|
eval_batch_size: int = 1
|
|
|
|
num_workers: int = 16
|
|
|
|
|
|
class MultiObjectDataset(Dataset):
|
|
def __init__(self, cfg: Any, split: str = "train") -> None:
|
|
super().__init__()
|
|
assert split in ["train", "val", "test"]
|
|
self.cfg: MultiObjectDataModuleConfig = cfg
|
|
|
|
self.all_scenes = _parse_scene_list(
|
|
self.cfg.scene_list, self.cfg.surface_root_dir
|
|
)
|
|
self.all_objects = _parse_object_list(self.cfg.object_list)
|
|
if len(self.all_scenes) != len(self.all_objects):
|
|
raise ValueError(
|
|
f"Number of scenes and objects must be the same, got {len(self.all_scenes)} scenes and {len(self.all_objects)} object lists."
|
|
)
|
|
|
|
self.all_images = _parse_scene_list(
|
|
self.cfg.scene_list, self.cfg.image_root_dir
|
|
)
|
|
|
|
self.split = split
|
|
self.indices = []
|
|
if self.split == "train" and self.cfg.train_indices is not None:
|
|
self.indices = (self.cfg.train_indices[0], self.cfg.train_indices[1])
|
|
elif self.split == "val" and self.cfg.val_indices is not None:
|
|
self.indices = (self.cfg.val_indices[0], self.cfg.val_indices[1])
|
|
elif self.split == "test" and self.cfg.test_indices is not None:
|
|
self.indices = (self.cfg.test_indices[0], self.cfg.test_indices[1])
|
|
else:
|
|
self.indices = (0, len(self.all_scenes))
|
|
|
|
repeat = self.cfg.repeat if self.split == "train" else 1
|
|
|
|
self.all_scenes = self.all_scenes[self.indices[0] : self.indices[1]] * repeat
|
|
self.all_objects = self.all_objects[self.indices[0] : self.indices[1]] * repeat
|
|
self.all_images = self.all_images[self.indices[0] : self.indices[1]] * repeat
|
|
|
|
if self.cfg.do_mix:
|
|
self.mix_all_scenes = _parse_scene_list(
|
|
self.cfg.mix_scene_list, self.cfg.mix_surface_root_dir
|
|
)[: self.cfg.mix_length]
|
|
self.mix_all_images = _parse_scene_list(
|
|
self.cfg.mix_scene_list, self.cfg.mix_image_root_dir
|
|
)[: self.cfg.mix_length]
|
|
|
|
def __len__(self):
|
|
return len(self.all_scenes)
|
|
|
|
def get_bg_color(self, bg_color):
|
|
if bg_color == "white":
|
|
bg_color = np.array([1.0, 1.0, 1.0], dtype=np.float32)
|
|
elif bg_color == "black":
|
|
bg_color = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
|
elif bg_color == "gray":
|
|
bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
|
|
elif bg_color == "random":
|
|
bg_color = np.random.rand(3)
|
|
elif bg_color == "random_gray":
|
|
bg_color = random.uniform(0.3, 0.7)
|
|
bg_color = np.array([bg_color] * 3, dtype=np.float32)
|
|
elif isinstance(bg_color, float):
|
|
bg_color = np.array([bg_color] * 3, dtype=np.float32)
|
|
elif isinstance(bg_color, list) or isinstance(bg_color, tuple):
|
|
bg_color = np.array(bg_color, dtype=np.float32)
|
|
else:
|
|
raise NotImplementedError
|
|
return bg_color
|
|
|
|
def load_surface(self, path, num_pc: int = 20480):
|
|
if path.endswith(".npy"):
|
|
data = np.load(path, allow_pickle=True).tolist()
|
|
surface = data["surface_points"] # Nx3
|
|
normal = data["surface_normals"] # Nx3
|
|
elif path.endswith(".obj") or path.endswith(".glb"):
|
|
import trimesh
|
|
|
|
n_surf_sample = 500000
|
|
scene = trimesh.load(path, process=False, force="scene")
|
|
meshes = []
|
|
|
|
for node_name in scene.graph.nodes_geometry:
|
|
geom_name = scene.graph[node_name][1]
|
|
geometry = scene.geometry[geom_name]
|
|
transform = scene.graph[node_name][0]
|
|
if isinstance(geometry, trimesh.Trimesh):
|
|
geometry.apply_transform(transform)
|
|
meshes.append(geometry)
|
|
mesh = trimesh.util.concatenate(meshes)
|
|
surface, face_indices = trimesh.sample.sample_surface(
|
|
mesh, n_surf_sample, sample_color=False
|
|
)
|
|
normal = mesh.face_normals[face_indices]
|
|
else:
|
|
raise NotImplementedError(f"Unsupported file format: {path}")
|
|
|
|
rng = np.random.default_rng()
|
|
ind = rng.choice(surface.shape[0], num_pc, replace=False)
|
|
surface = torch.FloatTensor(surface[ind])
|
|
normal = torch.FloatTensor(normal[ind])
|
|
surface = torch.cat([surface, normal], dim=-1)
|
|
|
|
return surface
|
|
|
|
def load_image(
|
|
self,
|
|
path,
|
|
height,
|
|
width,
|
|
background_color,
|
|
rescale: bool = False,
|
|
return_mask: bool = False,
|
|
remove_bg: bool = False,
|
|
idmap_path: Optional[str] = None,
|
|
):
|
|
image_pil = Image.open(path).resize((width, height))
|
|
image = torch.from_numpy(np.array(image_pil)).float() / 255.0
|
|
|
|
if image_pil.mode == "RGBA":
|
|
image_bg = image[:, :, :3] * image[:, :, 3:4] + background_color * (
|
|
1 - image[:, :, 3:4]
|
|
)
|
|
mask = (image[:, :, 3] > 0.5).float()
|
|
elif remove_bg and idmap_path is not None:
|
|
id_map = torch.from_numpy(
|
|
np.array(Image.open(idmap_path).resize((width, height), Image.NEAREST))
|
|
)
|
|
mask = (id_map > 0).float()
|
|
mask_ = mask.unsqueeze(-1).repeat(1, 1, 3)
|
|
image_bg = image * mask_ + background_color * (1 - mask_)
|
|
else:
|
|
image_bg = image
|
|
mask = torch.ones_like(image[:, :, 0]).float()
|
|
|
|
if rescale:
|
|
image_bg = image_bg * 2.0 - 1.0
|
|
if return_mask:
|
|
return image_bg, mask
|
|
return image_bg
|
|
|
|
def load_parts(
|
|
self,
|
|
rgb_path: str,
|
|
idmap_path: str,
|
|
indexes: List[int],
|
|
height: int,
|
|
width: int,
|
|
background_color: torch.Tensor,
|
|
skip_small_object: bool = False,
|
|
small_image_proportion: float = 0.005,
|
|
morph_perturb: bool = False, # Whether to apply morphological perturbation
|
|
max_kernel_size: int = 5,
|
|
p_dilation: float = 0.5,
|
|
):
|
|
rgb_image = self.load_image(rgb_path, height, width, background_color)
|
|
id_map = torch.from_numpy(
|
|
np.array(Image.open(idmap_path).resize((width, height), Image.NEAREST))
|
|
)
|
|
height, width, _ = rgb_image.shape
|
|
rgb_list, mask_list, success_list = [], [], []
|
|
|
|
for idx in indexes:
|
|
mask = (id_map == idx).float()
|
|
|
|
if (
|
|
skip_small_object
|
|
and mask.sum() <= small_image_proportion * height * width
|
|
):
|
|
success_list.append(False)
|
|
continue
|
|
|
|
if morph_perturb:
|
|
mask = random_morphological_transform(
|
|
mask, max_kernel_size=max_kernel_size, p_dilation=p_dilation
|
|
)
|
|
|
|
mask_3c = mask.unsqueeze(-1).repeat(1, 1, 3)
|
|
part_rgb = rgb_image * mask_3c + background_color * (1 - mask_3c)
|
|
rgb_list.append(part_rgb)
|
|
mask_list.append(mask)
|
|
success_list.append(True)
|
|
|
|
return rgb_list, mask_list, success_list
|
|
|
|
def crop_and_pad(self, rgbs, masks, height, width, padding_ratio=0.1):
|
|
cropped_rgbs, cropped_masks = [], []
|
|
|
|
for rgb, mask in zip(rgbs, masks):
|
|
rgb = rgb.permute(2, 0, 1)
|
|
|
|
# crop
|
|
coords = torch.nonzero(mask == 1)
|
|
y_min, x_min = coords.min(dim=0).values
|
|
y_max, x_max = coords.max(dim=0).values
|
|
|
|
cropped_rgb = rgb[:, y_min : y_max + 1, x_min : x_max + 1]
|
|
cropped_mask = mask[y_min : y_max + 1, x_min : x_max + 1]
|
|
|
|
h, w = cropped_rgb.shape[1:]
|
|
|
|
# padding
|
|
padding_size = [0, 0, 0, 0] # left, right, top, bottom
|
|
if w > h:
|
|
padding_size[2] = padding_size[3] = int((w - h) / 2)
|
|
h = w
|
|
else:
|
|
padding_size[0] = padding_size[1] = int((h - w) / 2)
|
|
w = h
|
|
|
|
padding_size = tuple([s + int(w * padding_ratio) for s in padding_size])
|
|
padded_rgb = F.pad(cropped_rgb, padding_size, mode="constant", value=1)
|
|
padded_mask = F.pad(cropped_mask, padding_size, mode="constant", value=0)
|
|
|
|
# resize
|
|
padded_rgb = F.interpolate(
|
|
padded_rgb.unsqueeze(0), (height, width), mode="bilinear"
|
|
)[0]
|
|
padded_mask = F.interpolate(
|
|
padded_mask.unsqueeze(0).unsqueeze(0), (height, width), mode="nearest"
|
|
)[0][0]
|
|
|
|
cropped_rgbs.append(padded_rgb)
|
|
cropped_masks.append(padded_mask)
|
|
|
|
return cropped_rgbs, cropped_masks
|
|
|
|
def _getitem_scene(self, index):
|
|
background_color = torch.as_tensor(self.get_bg_color(self.cfg.background_color))
|
|
|
|
# Surface
|
|
scene = self.all_scenes[index]
|
|
scene_objects = self.all_objects[index]
|
|
surfaces = []
|
|
for scene_object in scene_objects:
|
|
surface_path = os.path.join(
|
|
scene, f"{scene_object}.{self.cfg.surface_suffix}"
|
|
)
|
|
surface = self.load_surface(
|
|
surface_path, self.cfg.num_surface_samples_per_object
|
|
)
|
|
surfaces.append(surface)
|
|
surfaces = torch.stack(surfaces) # (num_instances, num_points, 6)
|
|
|
|
num_instances = surfaces.shape[0]
|
|
|
|
# Image
|
|
image_dir = self.all_images[index]
|
|
image_name = (
|
|
random.choice(self.cfg.image_names)
|
|
if self.split == "train"
|
|
else self.cfg.image_names[0]
|
|
)
|
|
image_prefix = (
|
|
[self.cfg.image_prefix]
|
|
if isinstance(self.cfg.image_prefix, str)
|
|
else self.cfg.image_prefix
|
|
)
|
|
image_prefix = (
|
|
random.choice(image_prefix) if self.split == "train" else image_prefix[0]
|
|
)
|
|
image_path = os.path.join(
|
|
image_dir, f"{image_prefix}_{image_name}.{self.cfg.image_suffix}"
|
|
)
|
|
idmap_path = (
|
|
os.path.join(
|
|
image_dir,
|
|
f"{self.cfg.idmap_prefix}_{image_name}.{self.cfg.idmap_suffix}",
|
|
)
|
|
.replace("_controlnet", "")
|
|
.replace("_inpaint", "")
|
|
)
|
|
|
|
# Load image and parts
|
|
rgb_scene = (
|
|
self.load_image(
|
|
image_path,
|
|
height=self.cfg.height,
|
|
width=self.cfg.width,
|
|
background_color=background_color,
|
|
remove_bg=self.cfg.remove_scene_bg,
|
|
idmap_path=idmap_path,
|
|
)
|
|
.unsqueeze(0)
|
|
.repeat(num_instances, 1, 1, 1)
|
|
.permute(0, 3, 1, 2)
|
|
)
|
|
rgbs, masks, success_list = self.load_parts(
|
|
image_path,
|
|
idmap_path,
|
|
list(range(1, num_instances + 1)),
|
|
self.cfg.height,
|
|
self.cfg.width,
|
|
background_color,
|
|
skip_small_object=self.cfg.skip_small_object,
|
|
small_image_proportion=self.cfg.small_image_proportion,
|
|
morph_perturb=self.cfg.morph_perturb,
|
|
max_kernel_size=self.cfg.max_kernel_size,
|
|
p_dilation=self.cfg.p_dilation,
|
|
)
|
|
if len(rgbs) == 0:
|
|
return self._getitem(random.randint(0, self.__len__() - 1))
|
|
|
|
rgb = torch.stack(rgbs).permute(0, 3, 1, 2)
|
|
mask = torch.stack(masks).unsqueeze(1)
|
|
|
|
# Update `surfaces`, `num_instances`, `rgb_scene` according to `success_list`
|
|
success_list = torch.tensor(success_list, dtype=torch.bool)
|
|
surfaces = surfaces[success_list]
|
|
num_instances = surfaces.shape[0]
|
|
rgb_scene = rgb_scene[success_list]
|
|
|
|
if self.cfg.max_num_instances is not None:
|
|
if num_instances > self.cfg.max_num_instances:
|
|
indices = torch.randperm(num_instances)[: self.cfg.max_num_instances]
|
|
surfaces = surfaces[indices]
|
|
rgb = rgb[indices]
|
|
mask = mask[indices]
|
|
rgb_scene = rgb_scene[indices]
|
|
num_instances = self.cfg.max_num_instances
|
|
|
|
# Scene id
|
|
scene_id = "-".join(image_dir.split("/")[-2:])
|
|
|
|
rv = {
|
|
"id": scene_id,
|
|
"num_instances": num_instances,
|
|
"surface": surfaces,
|
|
"rgb": rgb,
|
|
"mask": mask,
|
|
"rgb_scene": (
|
|
rgb_scene if self.cfg.use_scene_image else torch.zeros_like(rgb_scene)
|
|
),
|
|
}
|
|
|
|
if self.cfg.return_scene:
|
|
surface_scene = surfaces.view(-1, *surfaces.shape[2:])
|
|
rv.update({"surface_scene": surface_scene})
|
|
|
|
if self.cfg.return_crop_padded:
|
|
cropped_rgbs, cropped_masks = self.crop_and_pad(
|
|
rgbs, masks, self.cfg.height_crop_padded, self.cfg.width_crop_padded
|
|
)
|
|
cropped_rgb = torch.stack(cropped_rgbs)
|
|
cropped_mask = torch.stack(cropped_masks).unsqueeze(1)
|
|
|
|
rv.update(
|
|
{"rgb_crop_padded": cropped_rgb, "mask_crop_padded": cropped_mask}
|
|
)
|
|
|
|
if self.cfg.padding:
|
|
keys = [
|
|
"surface",
|
|
"rgb",
|
|
"mask",
|
|
"rgb_scene",
|
|
"rgb_crop_padded",
|
|
"mask_crop_padded",
|
|
]
|
|
if num_instances < self.cfg.num_instances_per_batch:
|
|
pad = self.cfg.num_instances_per_batch - num_instances
|
|
indices = torch.randint(
|
|
0, num_instances, (pad,), device=surfaces.device
|
|
)
|
|
updated_dict = {
|
|
k: torch.cat([v, v[indices]]) if k in keys else v
|
|
for k, v in rv.items()
|
|
}
|
|
else:
|
|
indices = torch.randperm(num_instances, device=surfaces.device)
|
|
indices = indices[: self.cfg.num_instances_per_batch]
|
|
updated_dict = {
|
|
k: v[indices] if k in keys else v for k, v in rv.items()
|
|
}
|
|
updated_dict.update({"num_instances": self.cfg.num_instances_per_batch})
|
|
rv = updated_dict
|
|
|
|
return rv
|
|
|
|
def _getitem_mix(self, index):
|
|
background_color = torch.as_tensor(self.get_bg_color(self.cfg.background_color))
|
|
|
|
surfaces, rgbs, masks = [], [], []
|
|
|
|
indexes = torch.randint(
|
|
0, len(self.mix_all_scenes), (self.cfg.num_instances_per_batch,)
|
|
)
|
|
for i in indexes:
|
|
scene = self.mix_all_scenes[i]
|
|
|
|
# Surface
|
|
surface_path = f"{scene}.{self.cfg.mix_surface_suffix}"
|
|
surface = self.load_surface(
|
|
surface_path, self.cfg.num_surface_samples_per_object
|
|
)
|
|
surfaces.append(surface)
|
|
|
|
# Image
|
|
image_dir = self.mix_all_images[i]
|
|
image_name = (
|
|
random.choice(self.cfg.mix_image_names)
|
|
if self.split == "train"
|
|
else self.cfg.mix_image_names[0]
|
|
)
|
|
image_path = os.path.join(
|
|
image_dir,
|
|
f"{self.cfg.mix_image_prefix}_{image_name}.{self.cfg.mix_image_suffix}",
|
|
)
|
|
|
|
rgb, mask = self.load_image(
|
|
image_path,
|
|
height=self.cfg.height,
|
|
width=self.cfg.width,
|
|
background_color=background_color,
|
|
return_mask=True,
|
|
)
|
|
rgbs.append(rgb)
|
|
masks.append(mask)
|
|
|
|
surfaces = torch.stack(surfaces) # (num_instances, num_points, 6)
|
|
rgb = torch.stack(rgbs).permute(0, 3, 1, 2)
|
|
mask = torch.stack(masks).unsqueeze(1)
|
|
|
|
# Scene ID
|
|
scene_id = self.mix_all_images[indexes[0]].split("/")[-1]
|
|
|
|
rv = {
|
|
"id": scene_id,
|
|
"num_instances": 1,
|
|
"surface": surfaces,
|
|
"rgb": rgb,
|
|
"mask": mask,
|
|
"rgb_scene": (
|
|
rgb if self.cfg.use_scene_image else torch.zeros_like(rgb)
|
|
), # here single object is the scene
|
|
}
|
|
|
|
return rv
|
|
|
|
def _getitem(self, index):
|
|
if (
|
|
self.split == "train"
|
|
and self.cfg.do_mix
|
|
and random.random() < self.cfg.do_mix_prob
|
|
):
|
|
return self._getitem_mix(index)
|
|
else:
|
|
return self._getitem_scene(index)
|
|
|
|
def __getitem__(self, index):
|
|
try:
|
|
return self._getitem(index)
|
|
except Exception as e:
|
|
print(f"Error in {self.all_scenes[index]}: {e}")
|
|
return self.__getitem__(random.randint(0, self.__len__() - 1))
|
|
|
|
def collate(self, batch):
|
|
batch = torch.utils.data.default_collate(batch)
|
|
pack = lambda t: t.view(-1, *t.shape[2:])
|
|
for k in batch.keys():
|
|
if k in [
|
|
"surface",
|
|
"rgb",
|
|
"mask",
|
|
"rgb_scene",
|
|
"rgb_crop_padded",
|
|
"mask_crop_padded",
|
|
]:
|
|
batch[k] = pack(batch[k])
|
|
|
|
batch["num_instances_per_batch"] = self.cfg.num_instances_per_batch
|
|
|
|
return batch
|
|
|
|
|
|
class MultiObjectDataModule(pl.LightningDataModule):
|
|
cfg: MultiObjectDataModuleConfig
|
|
|
|
def __init__(self, cfg: Optional[Union[dict, DictConfig]] = None) -> None:
|
|
super().__init__()
|
|
self.cfg = parse_structured(MultiObjectDataModuleConfig, cfg)
|
|
|
|
def setup(self, stage=None) -> None:
|
|
if stage in [None, "fit"]:
|
|
self.train_dataset = MultiObjectDataset(self.cfg, "train")
|
|
if stage in [None, "fit", "validate"]:
|
|
self.val_dataset = MultiObjectDataset(self.cfg, "val")
|
|
if stage in [None, "test", "predict"]:
|
|
self.test_dataset = MultiObjectDataset(self.cfg, "test")
|
|
|
|
def prepare_data(self):
|
|
pass
|
|
|
|
def train_dataloader(self) -> DataLoader:
|
|
return DataLoader(
|
|
self.train_dataset,
|
|
batch_size=self.cfg.batch_size,
|
|
num_workers=self.cfg.num_workers,
|
|
shuffle=True,
|
|
collate_fn=self.train_dataset.collate,
|
|
)
|
|
|
|
def val_dataloader(self) -> DataLoader:
|
|
return DataLoader(
|
|
self.val_dataset,
|
|
batch_size=self.cfg.eval_batch_size,
|
|
num_workers=self.cfg.num_workers,
|
|
shuffle=False,
|
|
collate_fn=self.val_dataset.collate,
|
|
)
|
|
|
|
def test_dataloader(self) -> DataLoader:
|
|
return DataLoader(
|
|
self.test_dataset,
|
|
batch_size=self.cfg.eval_batch_size,
|
|
num_workers=self.cfg.num_workers,
|
|
shuffle=False,
|
|
collate_fn=self.test_dataset.collate,
|
|
)
|
|
|
|
def predict_dataloader(self) -> DataLoader:
|
|
return self.test_dataloader()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import torchvision
|
|
from omegaconf import OmegaConf
|
|
|
|
config_file = "configs/scenediff/training.yaml"
|
|
data_cfg = OmegaConf.load(config_file)["data"]
|
|
cfg: MultiObjectDataModuleConfig = MultiObjectDataModuleConfig(**data_cfg)
|
|
data_module = MultiObjectDataModule(cfg)
|
|
data_module.setup()
|
|
|
|
for batch in data_module.test_dataloader():
|
|
print(batch["num_instances"])
|
|
|
|
for key in [
|
|
"rgb",
|
|
"mask",
|
|
"rgb_scene",
|
|
# "rgb_crop_padded",
|
|
# "mask_crop_padded",
|
|
]:
|
|
print(key, batch[key].shape, batch[key].min(), batch[key].max())
|
|
torchvision.utils.save_image(
|
|
batch[key], f"tmp/{key}.png", nrow=4, normalize=True
|
|
)
|
|
|
|
for key in ["rgb"]:
|
|
for i in range(batch[key].shape[0]):
|
|
torchvision.utils.save_image(
|
|
batch[key][i], f"tmp/{key}_{i}.png", normalize=True
|
|
)
|
|
|
|
break
|