from typing import * 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 import json import random import math import rembg import copy from pathlib import Path def _editverse_clean_encoded_base(clean_root: str, key: str, *, is_alpaca: bool) -> Path: """Root directory for one sample under 3Deditverse_data/encoded_ouput/.""" sub = "alpaca_encode" if is_alpaca else "flux_encode" return Path(clean_root) / "encoded_ouput" / sub / key def preprocess_condition_image(image: Image.Image, image_size: int, rembg_session=None, aug_bbox=None): """Preprocess editing condition images shared by 3DEditVerse and H3D.""" if np.array(image).shape[-1] != 4: if rembg_session is None: rembg_session = rembg.new_session('u2net') image = rembg.remove(image, session=rembg_session) alpha = np.array(image.getchannel(3)) if aug_bbox is None: 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_size_ratio = 1.2 aug_hsize = hsize * aug_size_ratio 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((image_size, 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) return image, aug_bbox, rembg_session class EditingStandardDatasetBase(Dataset): """ Base class for standard datasets. Args: roots (str): paths to the dataset """ def __init__(self, roots: str, json_file: str, dataset_num: Optional[int] = None, random_cond_gt: bool = False, # mixamo_data_path: Optional[str] = None, mixamo_data_repeat_ratio: float = 2.0, # flux_editing_image_path: Optional[str] = None, # flux_editing_edited_image_path: Optional[str] = None, # flux_editing_3d_path: Optional[str] = None, # alpaca_editing_image_path: Optional[str] = None, # alpaca_editing_3d_path: Optional[str] = None, # flux_editing_confidence_json_path: Optional[str] = None, # alpaca_editing_confidence_json_path: Optional[str] = None, # filter_info: Optional[dict] = None, adapt_simple_edit_data: bool = False, load_data_type: str = 'ss_latents', # voxel_num_json_path: Optional[str] = None, mixamo_root_path: Optional[str] = None, flux_edit_root_path: Optional[str] = None, alpaca_root_path: Optional[str] = None, dataset_info_path: Optional[str] = None, alpaca_confidence_json_path: Optional[str] = None, flux_edit_confidence_json_path: Optional[str] = None, filter_info: Optional[dict] = None, random_ori_edit: Optional[float] = None, flux_edit_alpaca_test_json_path: Optional[str] = None, mixamo_test_animation: Optional[List[str]] = None, simple_edit_data_if_filtered: bool = False, include_mixamo: bool = True, include_flux_edit: bool = True, include_alpaca: bool = True, max_samples: Optional[int] = None, editverse_clean_root: Optional[str] = None, ): super().__init__() self.roots = roots self.json_file = json_file self.dataset_num = dataset_num self.random_cond_gt = random_cond_gt # self.mixamo_data_path = mixamo_data_path self.mixamo_data_repeat_ratio = mixamo_data_repeat_ratio # self.flux_editing_image_path = flux_editing_image_path # self.flux_editing_edited_image_path = flux_editing_edited_image_path # self.flux_editing_3d_path = flux_editing_3d_path # self.alpaca_editing_image_path = alpaca_editing_image_path # self.alpaca_editing_3d_path = alpaca_editing_3d_path # self.flux_editing_confidence_json_path = flux_editing_confidence_json_path # self.alpaca_editing_confidence_json_path = alpaca_editing_confidence_json_path # self.filter_info = filter_info self.adapt_simple_edit_data = adapt_simple_edit_data self.load_data_type = load_data_type # self.voxel_num_json_path = voxel_num_json_path self.mixamo_root_path = mixamo_root_path self.flux_edit_root_path = flux_edit_root_path self.alpaca_root_path = alpaca_root_path self.dataset_info_path = dataset_info_path self.alpaca_confidence_json_path = alpaca_confidence_json_path self.flux_edit_confidence_json_path = flux_edit_confidence_json_path self.filter_info = filter_info self.random_ori_edit = random_ori_edit self.flux_edit_alpaca_test_json_path = flux_edit_alpaca_test_json_path self.mixamo_test_animation = mixamo_test_animation self.simple_edit_data_if_filtered = simple_edit_data_if_filtered self.include_mixamo = include_mixamo self.include_flux_edit = include_flux_edit self.include_alpaca = include_alpaca self.max_samples = max_samples self.editverse_clean_root = ( os.path.abspath(os.path.expanduser(editverse_clean_root)) if editverse_clean_root is not None else None ) if self.editverse_clean_root is not None: if adapt_simple_edit_data: raise ValueError( "adapt_simple_edit_data is incompatible with editverse_clean_root " "(clean encoded tree has no simple_edit_ss paths)." ) if simple_edit_data_if_filtered: raise ValueError( "simple_edit_data_if_filtered is incompatible with editverse_clean_root." ) assert self.load_data_type in ['ss_latents', 'latents'], f'Invalid load_data_type: {self.load_data_type}' if self.filter_info is not None: with open(self.flux_edit_confidence_json_path, 'r') as f: self.flux_edit_confidence_json = json.load(f) with open(self.alpaca_confidence_json_path, 'r') as f: self.alpaca_confidence_json = json.load(f) with open(self.dataset_info_path, 'r') as f: self.dataset_info = json.load(f) if self.random_ori_edit is not None: assert self.load_data_type == 'ss_latents', f'Invalid load_data_type: {self.load_data_type}' if self.flux_edit_alpaca_test_json_path is not None: with open(self.flux_edit_alpaca_test_json_path, 'r') as f: self.flux_edit_alpaca_test_json = json.load(f) self.flux_edit_alpaca_test_json = self.flux_edit_alpaca_test_json['alpaca'] + self.flux_edit_alpaca_test_json['flux_edit'] assert isinstance(self.flux_edit_alpaca_test_json, list) test_num = len(self.flux_edit_alpaca_test_json) self.flux_edit_alpaca_test_json = set(self.flux_edit_alpaca_test_json) assert test_num == len(self.flux_edit_alpaca_test_json) # v3 version # ---- prepare voxel num self.voxel_loads = [] # ---- load mixamo data self.mixamo_instances = self.dataset_info['mixamo'] if self.include_mixamo else {} print('Ori mixamo data num: ', sum(len(instances) for instances in self.mixamo_instances.values())) _mixamo_instances = [] for character_name, instances in self.mixamo_instances.items(): if self.mixamo_test_animation is not None and character_name in self.mixamo_test_animation: continue object_num = 0 for _ in range(math.ceil(self.mixamo_data_repeat_ratio)): random.shuffle(instances) i, j = 0, len(instances) - 1 while i < j: _mixamo_instances.append([ os.path.join(self.mixamo_root_path, instances[i]['ss_latents_path' if self.load_data_type == 'ss_latents' else 'latents_path']), os.path.join(self.mixamo_root_path, instances[j]['ss_latents_path' if self.load_data_type == 'ss_latents' else 'latents_path']), os.path.join(self.mixamo_root_path, instances[i]['img_path']), os.path.join(self.mixamo_root_path, instances[j]['img_path']), ]) self.voxel_loads.append((instances[i]['voxel_num'] + instances[j]['voxel_num']) // 2) object_num += 1 i += 1 j -= 1 if object_num >= len(instances) * self.mixamo_data_repeat_ratio: break self.mixamo_instances = _mixamo_instances print(f'Repeated Mixamo data num: {len(self.mixamo_instances)}') # ---- load flux editing data self.flux_editing_instances = [] remove_num = 0 test_num = 0 for key, value in (self.dataset_info['flux_edit'] if self.include_flux_edit else {}).items(): # load data ori_ss_latents_path = value['ori_ss_latents_path' if self.load_data_type == 'ss_latents' else 'ori_latents_path'] edited_ss_latents_path = value['edit_ss_latents_path' if self.load_data_type == 'ss_latents' else 'edit_latents_path'] # remove test data if self.flux_edit_alpaca_test_json_path is not None and key in self.flux_edit_alpaca_test_json: test_num += 1 continue # remove filtered data if self.filter_info is not None: choose_flag = True for filter_key, filter_range in filter_info.items(): min_val, max_val = filter_range if self.flux_edit_confidence_json[key][filter_key] < min_val or self.flux_edit_confidence_json[key][filter_key] > max_val: choose_flag = False break if not choose_flag: if not self.simple_edit_data_if_filtered: remove_num += 1 continue else: assert self.load_data_type == 'ss_latents' edited_ss_latents_path = edited_ss_latents_path.replace('edit_ss', 'simple_edit_ss').replace('data_v4', 'data_v2') assert os.path.exists(os.path.join(self.flux_edit_root_path, edited_ss_latents_path)), f'{edited_ss_latents_path} not found' # load data (release tree vs 3Deditverse_data encoded_ouput) if self.editverse_clean_root is None: if self.adapt_simple_edit_data: assert self.load_data_type == 'ss_latents' edited_ss_latents_path = edited_ss_latents_path.replace('edit_ss', 'simple_edit_ss').replace('data_v4', 'data_v2') assert os.path.exists(os.path.join(self.flux_edit_root_path, edited_ss_latents_path)), f'{edited_ss_latents_path} not found' self.flux_editing_instances.append([ os.path.join(self.flux_edit_root_path, ori_ss_latents_path), os.path.join(self.flux_edit_root_path, edited_ss_latents_path), os.path.join(self.flux_edit_root_path, value['ori_img_path']), os.path.join(self.flux_edit_root_path, value['edit_img_path']), ]) else: ori_latent, edit_latent = self._editverse_clean_latent_pair(key, is_alpaca=False) ori_img, edit_img = self._editverse_clean_cond_image_pair(key, value, is_alpaca=False) if not (os.path.isfile(ori_latent) and os.path.isfile(edit_latent)): remove_num += 1 continue if not (os.path.isfile(ori_img) and os.path.isfile(edit_img)): remove_num += 1 continue self.flux_editing_instances.append([ ori_latent, edit_latent, ori_img, edit_img, ]) self.voxel_loads.append((value['ori_voxel_num'] + value['edit_voxel_num']) // 2) print(f'Remove flux editing data (based on confidence) num: {remove_num}') print(f'Remove test flux editing data num: {test_num}') print(f'Flux editing data num: {len(self.flux_editing_instances)}') # ---- load alpaca editing data self.alpaca_editing_instances = [] remove_num = 0 test_num = 0 for key, value in (self.dataset_info['alpaca'] if self.include_alpaca else {}).items(): # load data ori_ss_latents_path = value['ori_ss_latents_path' if self.load_data_type == 'ss_latents' else 'ori_latents_path'] edited_ss_latents_path = value['edit_ss_latents_path' if self.load_data_type == 'ss_latents' else 'edit_latents_path'] # remove test data if self.flux_edit_alpaca_test_json_path is not None and key in self.flux_edit_alpaca_test_json: test_num += 1 continue # remove filtered data if self.filter_info is not None: choose_flag = True for filter_key, filter_range in filter_info.items(): min_val, max_val = filter_range if self.alpaca_confidence_json[key][filter_key] < min_val or self.alpaca_confidence_json[key][filter_key] > max_val: choose_flag = False break if not choose_flag: if not self.simple_edit_data_if_filtered: remove_num += 1 continue else: assert self.load_data_type == 'ss_latents' edited_ss_latents_path = edited_ss_latents_path.replace('edit_ss', 'simple_edit_ss').replace('data_V3', 'data_V2') assert os.path.exists(os.path.join(self.alpaca_root_path, edited_ss_latents_path)), f'{edited_ss_latents_path} not found' # load data (release tree vs 3Deditverse_data encoded_ouput) if self.editverse_clean_root is None: if self.adapt_simple_edit_data: assert self.load_data_type == 'ss_latents' edited_ss_latents_path = edited_ss_latents_path.replace('edit_ss', 'simple_edit_ss').replace('data_V3', 'data_V2') assert os.path.exists(os.path.join(self.alpaca_root_path, edited_ss_latents_path)), f'{edited_ss_latents_path} not found' self.alpaca_editing_instances.append([ os.path.join(self.alpaca_root_path, ori_ss_latents_path), os.path.join(self.alpaca_root_path, edited_ss_latents_path), os.path.join(self.alpaca_root_path, value['ori_img_path']), os.path.join(self.alpaca_root_path, value['edit_img_path']), ]) else: ori_latent, edit_latent = self._editverse_clean_latent_pair(key, is_alpaca=True) ori_img, edit_img = self._editverse_clean_cond_image_pair(key, value, is_alpaca=True) if not (os.path.isfile(ori_latent) and os.path.isfile(edit_latent)): remove_num += 1 continue if not (os.path.isfile(ori_img) and os.path.isfile(edit_img)): remove_num += 1 continue self.alpaca_editing_instances.append([ ori_latent, edit_latent, ori_img, edit_img, ]) self.voxel_loads.append((value['ori_voxel_num'] + value['edit_voxel_num']) // 2) print(f'Remove alpaca editing data (based on confidence) num: {remove_num}') print(f'Remove test alpaca editing data num: {test_num}') print(f'Alpaca editing data num: {len(self.alpaca_editing_instances)}') if self.max_samples is not None: original_num = len(self.mixamo_instances) + len(self.flux_editing_instances) + len(self.alpaca_editing_instances) remaining = self.max_samples self.mixamo_instances = self.mixamo_instances[:remaining] remaining -= len(self.mixamo_instances) self.flux_editing_instances = self.flux_editing_instances[:max(0, remaining)] remaining -= len(self.flux_editing_instances) self.alpaca_editing_instances = self.alpaca_editing_instances[:max(0, remaining)] self.voxel_loads = self.voxel_loads[:self.max_samples] print(f'Limit editing data num from {original_num} to {len(self.mixamo_instances) + len(self.flux_editing_instances) + len(self.alpaca_editing_instances)}') # v2 version '''# ---- load mixamo data self.mixamo_instances = defaultdict(list) if self.load_data_type == 'ss_latents': self.mixamo_ss_latents_path = os.path.join(self.mixamo_data_path, 'ss_latents', 'ss_enc_conv3d_16l8_fp16') else: self.mixamo_ss_latents_path = os.path.join(self.mixamo_data_path, 'latents', 'dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16') self.mixamo_renders_cond_path = os.path.join(self.mixamo_data_path, 'renders_cond') self.mixamo_ss_latents_list = os.listdir(self.mixamo_ss_latents_path) self.mixamo_ss_latents_list.sort() for ss_latent_name in self.mixamo_ss_latents_list: character_name = ss_latent_name.split('_')[0] cond_img_path = os.path.join(self.mixamo_data_path, 'renders_cond', ss_latent_name[:-4], '000.png') assert os.path.exists(cond_img_path), f'{cond_img_path} not found' self.mixamo_instances[character_name].append([ os.path.join(self.mixamo_ss_latents_path, ss_latent_name), os.path.join(self.mixamo_renders_cond_path, ss_latent_name[:-4], '000.png') ]) print('Ori mixamo data num: ', len(self.mixamo_ss_latents_list)) _mixamo_instances = [] for character_name, instances in self.mixamo_instances.items(): object_num = 0 for _ in range(math.ceil(self.mixamo_data_repeat_ratio)): random.shuffle(instances) i, j = 0, len(instances) - 1 while i < j: _mixamo_instances.append([ instances[i][0], instances[j][0], instances[i][1], instances[j][1] ]) object_num += 1 i += 1 j -= 1 if object_num >= len(instances) * self.mixamo_data_repeat_ratio: break self.mixamo_instances = _mixamo_instances print(f'Repeated Mixamo data num: {len(self.mixamo_instances)}') # ---- load flux editing data self.flux_editing_instances = [] self.flux_editing_3d_list = os.listdir(self.flux_editing_3d_path) self.flux_editing_3d_list.sort() remove_num = 0 for flux_editing_3d_name in self.flux_editing_3d_list: if self.load_data_type == 'ss_latents': ori_ss_latent_path = os.path.join(self.flux_editing_3d_path, flux_editing_3d_name, 'ori_ss_latents.npz') edited_ss_latent_path = os.path.join(self.flux_editing_3d_path, flux_editing_3d_name, 'edit_ss_latents.npz') else: ori_ss_latent_path = os.path.join(self.flux_editing_3d_path, flux_editing_3d_name, 'ori_latents.npz') edited_ss_latent_path = os.path.join(self.flux_editing_3d_path, flux_editing_3d_name, 'edit_latents.npz').replace('Flux_Step_3d_editing_data_v2', 'Flux_Step_3d_editing_data_v3') if not os.path.exists(edited_ss_latent_path): continue if not (os.path.exists(ori_ss_latent_path) and os.path.exists(edited_ss_latent_path)): continue if self.adapt_simple_edit_data: edited_ss_latent_path = edited_ss_latent_path.replace('edit_ss_latents.npz', 'simple_edit_ss_latents.npz') assert os.path.exists(edited_ss_latent_path), f'{edited_ss_latent_path} not found' if self.filter_info is not None and self.flux_editing_confidence_json_path is not None: choose_flag = True for filter_key, filter_range in filter_info.items(): min_val, max_val = filter_range if self.flux_editing_confidence_json[flux_editing_3d_name][filter_key] < min_val or self.flux_editing_confidence_json[flux_editing_3d_name][filter_key] > max_val: choose_flag = False break if not choose_flag: remove_num += 1 continue object_name, index = flux_editing_3d_name.split('_') ori_img_path = os.path.join(self.flux_editing_image_path, object_name, f'ori_{index}.png') # edited_img_path = os.path.join(self.flux_editing_image_path, object_name, f'edited_{index}.png') edit_img_list = os.listdir(os.path.join(self.flux_editing_edited_image_path, flux_editing_3d_name)) edit_img_list.sort() edited_img_path = os.path.join(self.flux_editing_edited_image_path, flux_editing_3d_name, edit_img_list[-1]) assert os.path.exists(ori_img_path), f'{ori_img_path} not found' assert os.path.exists(edited_img_path), f'{edited_img_path} not found' self.flux_editing_instances.append([ ori_ss_latent_path, edited_ss_latent_path, ori_img_path, edited_img_path, ]) print(f'Remove flux editing data (based on confidence) num: {remove_num}') print(f'Flux editing data num: {len(self.flux_editing_instances)}') # ---- load alpaca editing data self.alpaca_editing_instances = [] self.alpaca_editing_3d_list = os.listdir(self.alpaca_editing_3d_path) self.alpaca_editing_3d_list.sort() remove_num = 0 if self.load_data_type == 'ss_latents': for alpaca_editing_3d_name in self.alpaca_editing_3d_list: ori_ss_latent_path = os.path.join(self.alpaca_editing_3d_path, alpaca_editing_3d_name, 'ori_ss_latents.npz') edited_ss_latent_path = os.path.join(self.alpaca_editing_3d_path, alpaca_editing_3d_name, 'edit_ss_latents.npz') if not (os.path.exists(ori_ss_latent_path) and os.path.exists(edited_ss_latent_path)): continue if self.filter_info is not None and self.alpaca_editing_confidence_json_path is not None: choose_flag = True for filter_key, filter_range in filter_info.items(): min_val, max_val = filter_range if self.alpaca_editing_confidence_json[alpaca_editing_3d_name][filter_key] < min_val or self.alpaca_editing_confidence_json[alpaca_editing_3d_name][filter_key] > max_val: choose_flag = False break if not choose_flag: remove_num += 1 continue ori_img_path = os.path.join(self.alpaca_editing_image_path, alpaca_editing_3d_name, 'original.png') if not os.path.exists(ori_img_path): ori_img_path = ori_img_path.replace('open_data', 'open_data2') assert os.path.exists(ori_img_path), f'{ori_img_path} not found' edited_img_path = ori_img_path.replace('original.png', 'after_edited_Step1X-Edit.png') assert os.path.exists(edited_img_path), f'{edited_img_path} not found' self.alpaca_editing_instances.append([ ori_ss_latent_path, edited_ss_latent_path, ori_img_path, edited_img_path, ]) print(f'Remove alpaca editing data (based on confidence) num: {remove_num}') print(f'Alpaca editing data num: {len(self.alpaca_editing_instances)}')''' print('**********************') print(f'Total editing data num: {len(self.mixamo_instances) + len(self.flux_editing_instances) + len(self.alpaca_editing_instances)}, including:') print(f' - Mixamo: {len(self.mixamo_instances)}') print(f' - Flux editing: {len(self.flux_editing_instances)}') print(f' - Alpaca editing: {len(self.alpaca_editing_instances)}') print('**********************') # v1 version '''with open(json_file, 'r') as f: json_info = json.load(f) json_info = json_info[:dataset_num] print('Plan to load {} instances'.format(len(json_info))) self.instances = [] for info in json_info: glb1_path = info['model1_path'] merged_glb_path = info['merged_glb_path'] glb1_sha256 = glb1_path.split('/')[-1].split('.')[0] merged_glb_sha256 = merged_glb_path.split('/')[-1].split('.')[0] exist_glb1 = self.exist_instance(self.roots, glb1_sha256) exist_merged_glb = self.exist_instance(self.roots, merged_glb_sha256) if exist_glb1 and exist_merged_glb: self.instances.append([self.roots, glb1_sha256, merged_glb_sha256]) print('Actually load {} instances'.format(len(self.instances)))''' # ori version '''self.metadata = pd.DataFrame() self._stats = {} 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) self._stats[key].update(stats) self.instances.extend([(root, sha256) for sha256 in metadata['sha256'].values]) metadata.set_index('sha256', inplace=True) self.metadata = pd.concat([self.metadata, metadata])''' def _editverse_clean_latent_pair(self, key: str, *, is_alpaca: bool) -> Tuple[str, str]: base = _editverse_clean_encoded_base(self.editverse_clean_root, key, is_alpaca=is_alpaca) if self.load_data_type == 'ss_latents': return str(base / 'ori' / 'ss_latent.npz'), str(base / 'edit' / 'ss_latent.npz') return str(base / 'ori' / 'slat_latent.npz'), str(base / 'edit' / 'slat_latent.npz') def _editverse_clean_cond_image_pair(self, key: str, value: dict, *, is_alpaca: bool) -> Tuple[str, str]: branch = 'alpaca' if is_alpaca else 'flux' cond_base = Path(self.editverse_clean_root) / 'condition-image' / branch / key ori_candidates = [ cond_base / os.path.basename(value['ori_img_path']), cond_base / 'ori_image.png', ] edit_candidates = [ cond_base / os.path.basename(value['edit_img_path']), cond_base / ('after_edited_Flux.png' if is_alpaca else 'edited_0.png'), ] ori = next((path for path in ori_candidates if path.is_file()), ori_candidates[0]) edit = next((path for path in edit_candidates if path.is_file()), edit_candidates[0]) return str(ori), str(edit) def exist_instance(self, root, sha256): cond1 = os.path.exists(os.path.join(root, 'ss_latents', 'ss_enc_conv3d_16l8_fp16', f'{sha256}.npz')) cond2 = os.path.exists(os.path.join(root, 'renders_cond', sha256, 'transforms.json')) return cond1 and cond2 @abstractmethod def filter_metadata(self, metadata: pd.DataFrame) -> Tuple[pd.DataFrame, Dict[str, int]]: pass # @abstractmethod # def get_instance(self, root: str, glb1_sha256: str, merged_glb_sha256: str) -> Dict[str, Any]: # pass @abstractmethod def get_instance(self, ori_ss_latent_path: str, edited_ss_latent_path: str, ori_img_path: str, edited_img_path: str) -> Dict[str, Any]: pass def __len__(self): # return len(self.instances) return len(self.mixamo_instances) + len(self.flux_editing_instances) + len(self.alpaca_editing_instances) def get_data_path_from_index(self, index) -> Dict[str, Any]: if index < len(self.mixamo_instances): data = self.mixamo_instances[index] elif index < len(self.mixamo_instances) + len(self.flux_editing_instances): data = self.flux_editing_instances[index - len(self.mixamo_instances)] else: data = self.alpaca_editing_instances[index - len(self.mixamo_instances) - len(self.flux_editing_instances)] return data def __getitem__(self, index) -> Dict[str, Any]: # try: data_path = self.get_data_path_from_index(index) if self.random_ori_edit is not None and np.random.rand() < self.random_ori_edit: another_index = np.random.randint(0, len(self)) another_data_path = self.get_data_path_from_index(another_index) data_path = [another_data_path[0], data_path[1], another_data_path[2], data_path[3]] return self.get_instance(*data_path) # except Exception as e: # print(e) # return self.__getitem__(np.random.randint(0, len(self))) def __str__(self): lines = [] lines.append(self.__class__.__name__) # v2 version lines.append(f' - Total instances: {len(self)}') lines.append(f' - Sources:') lines.append(f' - Mixamo: {len(self.mixamo_instances)}') lines.append(f' - Flux editing: {len(self.flux_editing_instances)}') lines.append(f' - Alpaca editing: {len(self.alpaca_editing_instances)}') return '\n'.join(lines) # v1 version '''lines.append(f' - Total instances: {len(self)}') lines.append(f' - Sources:') for root, glb1_sha256, merged_glb_sha256 in self.instances: lines.append(f' - {root}:') lines.append(f' - {glb1_sha256}:') lines.append(f' - {merged_glb_sha256}:') return '\n'.join(lines)''' class EditingImageConditionedMixin: def __init__(self, roots, *, image_size=518, **kwargs): self.image_size = image_size super().__init__(roots, **kwargs) # lazy init rembg session self.rembg_session = None def __deepcopy__(self, memo): """ 自定义深度复制逻辑。 """ # print("正在调用自定义的 __deepcopy__ ...") # 1. 创建一个新的类实例，但不调用 __init__，避免重复创建 session cls = self.__class__ result = cls.__new__(cls) # memo 是一个字典，用来处理循环引用，这是实现 __deepcopy__ 的标准做法 memo[id(self)] = result # 2. 遍历 self 的 __dict__，深度复制所有其他属性 for k, v in self.__dict__.items(): if k == 'rembg_session': # 跳过 rembg_session 的复制 # print(f"跳过复制 'rembg_session' 属性") continue setattr(result, k, copy.deepcopy(v, memo)) # 3. 在新的副本中，手动创建一个全新的 rembg_session # print("在新的副本中创建一个全新的 rembg_session...") result.rembg_session = rembg.new_session('u2net') return result # def filter_metadata(self, metadata): # metadata, stats = super().filter_metadata(metadata) # metadata = metadata[metadata[f'cond_rendered']] # stats['Cond rendered'] = len(metadata) # return metadata, stats def _get_cond_image(self, image_path, aug_bbox=None): # v1 version '''image_root = os.path.join(root, 'renders_cond', sha256) 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) if view_idx is None else view_idx metadata = metadata['frames'][view] image_path = os.path.join(image_root, metadata['file_path'])''' image = Image.open(image_path) image, aug_bbox, self.rembg_session = preprocess_condition_image( image, self.image_size, self.rembg_session, aug_bbox ) # return image, view, aug_bbox return image, aug_bbox def get_instance(self, ori_ss_latent_path, edited_ss_latent_path, ori_img_path, edited_img_path): pack = super().get_instance(ori_ss_latent_path, edited_ss_latent_path, ori_img_path, edited_img_path) ori_cond_img, _ = self._get_cond_image(ori_img_path) edited_cond_img, _ = self._get_cond_image(edited_img_path) pack['ori_cond_img'] = ori_cond_img pack['edited_cond_img'] = edited_cond_img if self.random_cond_gt and np.random.rand() < 0.5: pack['ori_cond_img'], pack['edited_cond_img'] = pack['edited_cond_img'], pack['ori_cond_img'] pack['ori_ss_latent'], pack['edited_ss_latent'] = pack['edited_ss_latent'], pack['ori_ss_latent'] pack.setdefault('edit_type', 'editverse') return pack