diff --git a/gen_3d_modules/LGM/core/options.py b/gen_3d_modules/LGM/core/options.py index ba6e8f1..0136e56 100644 --- a/gen_3d_modules/LGM/core/options.py +++ b/gen_3d_modules/LGM/core/options.py @@ -55,10 +55,6 @@ class Options: mixed_precision: str = 'bf16' # learning rate lr: float = 4e-4 - # augmentation prob for grid distortion - prob_grid_distortion: float = 0.5 - # augmentation prob for camera jitter - prob_cam_jitter: float = 0.5 ### testing # test image path diff --git a/gen_3d_modules/LGM/core/provider_objaverse.py b/gen_3d_modules/LGM/core/provider_objaverse.py index 310eaee..3b94bca 100644 --- a/gen_3d_modules/LGM/core/provider_objaverse.py +++ b/gen_3d_modules/LGM/core/provider_objaverse.py @@ -1,17 +1,14 @@ import os import cv2 -import random import numpy as np import torch -import torch.nn as nn import torch.nn.functional as F import torchvision.transforms.functional as TF from torch.utils.data import Dataset -import kiui from LGM.core.options import Options -from LGM.core.utils import get_rays, grid_distortion, orbit_camera_jitter +from LGM.core.utils import get_rays IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406) IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225) @@ -130,15 +127,6 @@ class ObjaverseDataset(Dataset): images_input = F.interpolate(images[:self.opt.num_input_views].clone(), size=(self.opt.input_size, self.opt.input_size), mode='bilinear', align_corners=False) # [V, C, H, W] cam_poses_input = cam_poses[:self.opt.num_input_views].clone() - # data augmentation - if self.training: - # apply random grid distortion to simulate 3D inconsistency - if random.random() < self.opt.prob_grid_distortion: - images_input[1:] = grid_distortion(images_input[1:]) - # apply camera jittering (only to input!) - if random.random() < self.opt.prob_cam_jitter: - cam_poses_input[1:] = orbit_camera_jitter(cam_poses_input[1:]) - images_input = TF.normalize(images_input, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD) # resize render ground-truth images, range still in [0, 1] diff --git a/gen_3d_modules/LGM/core/utils.py b/gen_3d_modules/LGM/core/utils.py index 655f921..109fc79 100644 --- a/gen_3d_modules/LGM/core/utils.py +++ b/gen_3d_modules/LGM/core/utils.py @@ -62,53 +62,6 @@ def orbit_camera_jitter(poses, strength=0.1): return new_poses -def grid_distortion(images, strength=0.5): - # images: [B, C, H, W] - # num_steps: int, grid resolution for distortion - # strength: float in [0, 1], strength of distortion - - B, C, H, W = images.shape - - num_steps = np.random.randint(8, 17) - grid_steps = torch.linspace(-1, 1, num_steps) - - # have to loop batch... - grids = [] - for b in range(B): - # construct displacement - x_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive - x_steps = (x_steps + strength * (torch.rand_like(x_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb - x_steps = (x_steps * W).long() # [num_steps] - x_steps[0] = 0 - x_steps[-1] = W - xs = [] - for i in range(num_steps - 1): - xs.append(torch.linspace(grid_steps[i], grid_steps[i + 1], x_steps[i + 1] - x_steps[i])) - xs = torch.cat(xs, dim=0) # [W] - - y_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive - y_steps = (y_steps + strength * (torch.rand_like(y_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb - y_steps = (y_steps * H).long() # [num_steps] - y_steps[0] = 0 - y_steps[-1] = H - ys = [] - for i in range(num_steps - 1): - ys.append(torch.linspace(grid_steps[i], grid_steps[i + 1], y_steps[i + 1] - y_steps[i])) - ys = torch.cat(ys, dim=0) # [H] - - # construct grid - grid_x, grid_y = torch.meshgrid(xs, ys, indexing='xy') # [H, W] - grid = torch.stack([grid_x, grid_y], dim=-1) # [H, W, 2] - - grids.append(grid) - - grids = torch.stack(grids, dim=0).to(images.device) # [B, H, W, 2] - - # grid sample - images = F.grid_sample(images, grids, align_corners=False) - - return images - class WarningFilter(logging.Filter): def filter(self, record): if record.levelno == logging.WARNING: diff --git a/requirements.txt b/requirements.txt index 2ad87cc..c764268 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,10 +9,8 @@ scipy kornia opencv-python pillow -roma nerfacc>=0.5.3 PyMCubes -scikit-learn # for use ML models diffusers>=0.26.1