tencentarc--pixal3d
wehub-resource-sync
350 行
15 KiB
Python
350 行
15 KiB
Python
from typing import *
|
|
import json
|
|
from abc import abstractmethod
|
|
import os
|
|
import json
|
|
import torch
|
|
import numpy as np
|
|
import pandas as pd
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset
|
|
|
|
|
|
class StandardDatasetBase(Dataset):
|
|
"""
|
|
Base class for standard datasets.
|
|
|
|
Args:
|
|
roots (str): paths to the dataset
|
|
skip_list (str, optional): path to a file containing sha256 hashes to skip (one per line)
|
|
Format: "dataset/sha256" (e.g., "ABO/6a79dbb5...")
|
|
skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
|
|
(e.g., ["texverse"] for datasets without aesthetic_score)
|
|
"""
|
|
|
|
def __init__(self,
|
|
roots: str,
|
|
skip_list: Optional[str] = None,
|
|
skip_aesthetic_score_datasets: Optional[List[str]] = None,
|
|
):
|
|
super().__init__()
|
|
|
|
# Datasets to skip aesthetic score check
|
|
self.skip_aesthetic_score_datasets = set(skip_aesthetic_score_datasets or [])
|
|
|
|
# Load skip list if provided
|
|
self.skip_set = set()
|
|
if skip_list is not None and os.path.exists(skip_list):
|
|
with open(skip_list, 'r') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line and not line.startswith('#'):
|
|
self.skip_set.add(line)
|
|
print(f'Loaded {len(self.skip_set)} items from skip_list: {skip_list}')
|
|
|
|
try:
|
|
self.roots = json.loads(roots)
|
|
root_type = 'obj'
|
|
except:
|
|
self.roots = roots.split(',')
|
|
root_type = 'list'
|
|
self.instances = []
|
|
self.metadata = pd.DataFrame()
|
|
|
|
self._stats = {}
|
|
if root_type == 'obj':
|
|
for key, root in self.roots.items():
|
|
self._stats[key] = {}
|
|
metadata = pd.DataFrame(columns=['sha256']).set_index('sha256')
|
|
|
|
# Only merge key fields from ss_latent and render_cond
|
|
# Exclude base, because cond_rendered=False in base/metadata.csv would incorrectly overwrite real values
|
|
for sub_key, r in root.items():
|
|
if sub_key == 'base':
|
|
continue # Skip base directory
|
|
metadata_file = os.path.join(r, 'metadata.csv')
|
|
if os.path.exists(metadata_file):
|
|
metadata = metadata.combine_first(pd.read_csv(metadata_file).set_index('sha256'))
|
|
|
|
# Read aesthetic_score separately from base (avoid reading other potentially conflicting columns)
|
|
if 'base' in root:
|
|
base_metadata_file = os.path.join(root['base'], 'metadata.csv')
|
|
if os.path.exists(base_metadata_file):
|
|
base_df = pd.read_csv(base_metadata_file).set_index('sha256')
|
|
if 'aesthetic_score' in base_df.columns and 'aesthetic_score' not in metadata.columns:
|
|
metadata['aesthetic_score'] = base_df['aesthetic_score']
|
|
|
|
self._stats[key]['Total'] = len(metadata)
|
|
metadata, stats = self.filter_metadata(metadata, dataset_name=key)
|
|
self._stats[key].update(stats)
|
|
|
|
# Filter out items in skip_list
|
|
skipped_count = 0
|
|
for sha256 in metadata.index.values:
|
|
skip_key = f'{key}/{sha256}'
|
|
if skip_key in self.skip_set:
|
|
skipped_count += 1
|
|
else:
|
|
self.instances.append((root, sha256, key))
|
|
if skipped_count > 0:
|
|
self._stats[key]['Skipped (skip_list)'] = skipped_count
|
|
self._stats[key]['After skip_list'] = len(metadata) - skipped_count
|
|
|
|
self.metadata = pd.concat([self.metadata, metadata])
|
|
else:
|
|
for root in self.roots:
|
|
key = os.path.basename(root)
|
|
self._stats[key] = {}
|
|
metadata = pd.read_csv(os.path.join(root, 'metadata.csv'))
|
|
self._stats[key]['Total'] = len(metadata)
|
|
metadata, stats = self.filter_metadata(metadata, dataset_name=key)
|
|
self._stats[key].update(stats)
|
|
|
|
# Filter out items in skip_list
|
|
skipped_count = 0
|
|
for sha256 in metadata['sha256'].values:
|
|
skip_key = f'{key}/{sha256}'
|
|
if skip_key in self.skip_set:
|
|
skipped_count += 1
|
|
else:
|
|
self.instances.append((root, sha256, key))
|
|
if skipped_count > 0:
|
|
self._stats[key]['Skipped (skip_list)'] = skipped_count
|
|
self._stats[key]['After skip_list'] = len(metadata) - skipped_count
|
|
metadata.set_index('sha256', inplace=True)
|
|
self.metadata = pd.concat([self.metadata, metadata])
|
|
|
|
@abstractmethod
|
|
def filter_metadata(self, metadata: pd.DataFrame, dataset_name: str = None) -> Tuple[pd.DataFrame, Dict[str, int]]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_instance(self, root, instance: str) -> Dict[str, Any]:
|
|
pass
|
|
|
|
def __len__(self):
|
|
return len(self.instances)
|
|
|
|
def __getitem__(self, index) -> Dict[str, Any]:
|
|
try:
|
|
root, instance, dataset_name = self.instances[index]
|
|
pack = self.get_instance(root, instance)
|
|
pack['_dataset_name'] = dataset_name
|
|
pack['_sha256'] = instance
|
|
return pack
|
|
except Exception as e:
|
|
print(f'Error loading {self.instances[index][1]}: {e}')
|
|
return self.__getitem__(np.random.randint(0, len(self)))
|
|
|
|
def __str__(self):
|
|
lines = []
|
|
lines.append(self.__class__.__name__)
|
|
lines.append(f' - Total instances: {len(self)}')
|
|
lines.append(f' - Sources:')
|
|
for key, stats in self._stats.items():
|
|
lines.append(f' - {key}:')
|
|
for k, v in stats.items():
|
|
lines.append(f' - {k}: {v}')
|
|
return '\n'.join(lines)
|
|
|
|
|
|
class ImageConditionedMixin:
|
|
def __init__(self, roots, *, image_size=518, **kwargs):
|
|
self.image_size = image_size
|
|
super().__init__(roots, **kwargs)
|
|
|
|
def filter_metadata(self, metadata, dataset_name=None):
|
|
metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
|
|
metadata = metadata[metadata['cond_rendered'].notna()]
|
|
stats['Cond rendered'] = len(metadata)
|
|
return metadata, stats
|
|
|
|
def get_instance(self, root, instance):
|
|
pack = super().get_instance(root, instance)
|
|
|
|
image_root = os.path.join(root['render_cond'], instance)
|
|
with open(os.path.join(image_root, 'transforms.json')) as f:
|
|
metadata = json.load(f)
|
|
n_views = len(metadata['frames'])
|
|
view = np.random.randint(n_views)
|
|
metadata = metadata['frames'][view]
|
|
|
|
image_path = os.path.join(image_root, metadata['file_path'])
|
|
image = Image.open(image_path)
|
|
|
|
alpha = np.array(image.getchannel(3))
|
|
bbox = np.array(alpha).nonzero()
|
|
bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
|
|
center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
|
|
hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
|
|
aug_hsize = hsize
|
|
aug_center_offset = [0, 0]
|
|
aug_center = [center[0] + aug_center_offset[0], center[1] + aug_center_offset[1]]
|
|
aug_bbox = [int(aug_center[0] - aug_hsize), int(aug_center[1] - aug_hsize), int(aug_center[0] + aug_hsize), int(aug_center[1] + aug_hsize)]
|
|
image = image.crop(aug_bbox)
|
|
|
|
image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
|
alpha = image.getchannel(3)
|
|
image = image.convert('RGB')
|
|
image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
|
|
alpha = torch.tensor(np.array(alpha)).float() / 255.0
|
|
image = image * alpha.unsqueeze(0)
|
|
pack['cond'] = image
|
|
|
|
return pack
|
|
|
|
|
|
class ViewImageConditionedMixin:
|
|
"""
|
|
Mixin for view-based image-conditioned datasets.
|
|
|
|
This mixin is designed for datasets where ss_latent is stored per-view (view{XX}.npz),
|
|
and needs to load the corresponding view image and scale from view{XX}_scale.json.
|
|
|
|
Args:
|
|
image_size: Target image size
|
|
load_camera_info: Whether to load camera information for view-aligned conditioning
|
|
"""
|
|
def __init__(self, roots, *, image_size=518, load_camera_info=False, **kwargs):
|
|
self.image_size = image_size
|
|
# self.load_camera_info = load_camera_info
|
|
super().__init__(roots, **kwargs)
|
|
|
|
def filter_metadata(self, metadata, dataset_name=None):
|
|
metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
|
|
metadata = metadata[metadata['cond_rendered'].notna()]
|
|
stats['Cond rendered'] = len(metadata)
|
|
return metadata, stats
|
|
|
|
def get_instance(self, root, instance):
|
|
"""
|
|
Get instance with view-aligned image and camera info.
|
|
|
|
Expects parent class to set:
|
|
- pack['x_0']: the latent tensor
|
|
- self._current_view_idx: the selected view index
|
|
- self._current_latent_dir: the latent directory path
|
|
"""
|
|
pack = super().get_instance(root, instance)
|
|
|
|
# Get view_idx from parent class (set by SparseStructureLatentView)
|
|
if not hasattr(self, '_current_view_idx'):
|
|
raise RuntimeError("Parent class must set '_current_view_idx' before calling ViewImageConditionedMixin.get_instance")
|
|
if not hasattr(self, '_current_latent_dir'):
|
|
raise RuntimeError("Parent class must set '_current_latent_dir' before calling ViewImageConditionedMixin.get_instance")
|
|
view_idx = self._current_view_idx
|
|
latent_dir = self._current_latent_dir
|
|
|
|
# Load image metadata
|
|
image_root = os.path.join(root['render_cond'], instance)
|
|
with open(os.path.join(image_root, 'transforms.json')) as f:
|
|
metadata = json.load(f)
|
|
|
|
# Load corresponding image for this view
|
|
frame_metadata = metadata['frames'][view_idx]
|
|
image_path = os.path.join(image_root, frame_metadata['file_path'])
|
|
image = Image.open(image_path)
|
|
|
|
image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
|
alpha = image.getchannel(3)
|
|
image = image.convert('RGB')
|
|
image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
|
|
alpha = torch.tensor(np.array(alpha)).float() / 255.0
|
|
image = image * alpha.unsqueeze(0)
|
|
pack['cond'] = image
|
|
|
|
# Load camera info if requested
|
|
|
|
# camera_angle_x: check frame first, then root metadata
|
|
if 'camera_angle_x' in frame_metadata:
|
|
camera_angle_x = float(frame_metadata['camera_angle_x'])
|
|
elif 'camera_angle_x' in metadata:
|
|
camera_angle_x = float(metadata['camera_angle_x'])
|
|
else:
|
|
raise KeyError(f"'camera_angle_x' not found in transforms.json for {instance}")
|
|
pack['camera_angle_x'] = torch.tensor(camera_angle_x, dtype=torch.float32)
|
|
|
|
# transform_matrix
|
|
if 'transform_matrix' not in frame_metadata:
|
|
raise KeyError(f"'transform_matrix' not found in frame {view_idx} for {instance}")
|
|
transform_matrix = torch.tensor(frame_metadata['transform_matrix'], dtype=torch.float32)
|
|
distance = torch.norm(transform_matrix[:3, 3]).item()
|
|
|
|
pack['camera_distance'] = torch.tensor(distance, dtype=torch.float32)
|
|
# NOTE: Do NOT pass transform_matrix to ProjGrid.
|
|
# shape_latent space objects are already rotated to front-view by transform_mesh,
|
|
# so ProjGrid should use the default front_view_transform_matrix + distance.
|
|
# pack['transform_matrix'] = transform_matrix
|
|
|
|
# Load mesh_scale from ss_latent directory's view{XX}_scale.json
|
|
scale_json_path = os.path.join(latent_dir, f'view{view_idx:02d}_scale.json')
|
|
if not os.path.exists(scale_json_path):
|
|
raise FileNotFoundError(f"Scale file not found: {scale_json_path}")
|
|
with open(scale_json_path) as f:
|
|
scale_data = json.load(f)
|
|
if 'total_scale' not in scale_data:
|
|
raise KeyError(f"'total_scale' not found in {scale_json_path}")
|
|
pack['mesh_scale'] = torch.tensor(float(scale_data['total_scale']), dtype=torch.float32)
|
|
|
|
return pack
|
|
|
|
|
|
class MultiImageConditionedMixin:
|
|
def __init__(self, roots, *, image_size=518, max_image_cond_view = 4, **kwargs):
|
|
self.image_size = image_size
|
|
self.max_image_cond_view = max_image_cond_view
|
|
super().__init__(roots, **kwargs)
|
|
|
|
def filter_metadata(self, metadata, dataset_name=None):
|
|
metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
|
|
metadata = metadata[metadata['cond_rendered'].notna()]
|
|
stats['Cond rendered'] = len(metadata)
|
|
return metadata, stats
|
|
|
|
def get_instance(self, root, instance):
|
|
pack = super().get_instance(root, instance)
|
|
|
|
image_root = os.path.join(root['render_cond'], instance)
|
|
with open(os.path.join(image_root, 'transforms.json')) as f:
|
|
metadata = json.load(f)
|
|
|
|
n_views = len(metadata['frames'])
|
|
n_sample_views = np.random.randint(1, self.max_image_cond_view+1)
|
|
|
|
assert n_views >= n_sample_views, f'Not enough views to sample {n_sample_views} unique images.'
|
|
|
|
sampled_views = np.random.choice(n_views, size=n_sample_views, replace=False)
|
|
|
|
cond_images = []
|
|
for v in sampled_views:
|
|
frame_info = metadata['frames'][v]
|
|
image_path = os.path.join(image_root, frame_info['file_path'])
|
|
image = Image.open(image_path)
|
|
|
|
alpha = np.array(image.getchannel(3))
|
|
bbox = np.array(alpha).nonzero()
|
|
bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
|
|
center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
|
|
hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
|
|
aug_hsize = hsize
|
|
aug_center = center
|
|
aug_bbox = [
|
|
int(aug_center[0] - aug_hsize),
|
|
int(aug_center[1] - aug_hsize),
|
|
int(aug_center[0] + aug_hsize),
|
|
int(aug_center[1] + aug_hsize),
|
|
]
|
|
|
|
img = image.crop(aug_bbox)
|
|
img = img.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
|
alpha = img.getchannel(3)
|
|
img = img.convert('RGB')
|
|
img = torch.tensor(np.array(img)).permute(2, 0, 1).float() / 255.0
|
|
alpha = torch.tensor(np.array(alpha)).float() / 255.0
|
|
img = img * alpha.unsqueeze(0)
|
|
|
|
cond_images.append(img)
|
|
|
|
pack['cond'] = [torch.stack(cond_images, dim=0)] # (V,3,H,W)
|
|
return pack
|