content: # -*- coding: utf-8 -*-
  
  # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
  # except for the third-party components listed below.
  # Hunyuan 3D does not impose any additional limitations beyond what is outlined
  # in the repsective licenses of these third-party components.
  # Users must comply with all terms and conditions of original licenses of these third-party
  # components and must ensure that the usage of the third party components adheres to
  # all relevant laws and regulations.
  
  # For avoidance of doubts, Hunyuan 3D means the large language models and
  # their software and algorithms, including trained model weights, parameters (including
  # optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
  # fine-tuning enabling code and other elements of the foregoing made publicly available
  # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
  
  
  import os
  import io
  import sys
  import time
  import random
  import traceback
  from typing import Optional, Union, List, Tuple, Dict
  
  import json
  import glob
  import cv2
  import numpy as np
  import trimesh
  
  import torch
  import torchvision.transforms as transforms
  from pytorch_lightning import LightningDataModule
  from pytorch_lightning.utilities import rank_zero_info
  
  from .utils import worker_init_fn, pytorch_worker_seed, make_seed
  
  
  class ResampledShards(torch.utils.data.dataset.IterableDataset):
      def __init__(self, datalist, nshards=sys.maxsize, worker_seed=None, deterministic=False):
          super().__init__()
          self.datalist = datalist
          self.nshards = nshards
          # If no worker_seed provided, use pytorch_worker_seed function; else use given seed
          self.worker_seed = pytorch_worker_seed if worker_seed is None else worker_seed
          self.deterministic = deterministic
          self.epoch = -1
  
      def __iter__(self):
          self.epoch += 1
          if self.deterministic:
              seed = make_seed(self.worker_seed(), self.epoch)
          else:
              seed = make_seed(self.worker_seed(), self.epoch, 
                               os.getpid(), time.time_ns(), os.urandom(4))
          self.rng = random.Random(seed)
          for _ in range(self.nshards):
              index = self.rng.randint(0, len(self.datalist) - 1)
              yield self.datalist[index]
  
              
  def read_npz(data):
      # Load a numpy .npz file from a file path or file-like object
      # The commented line shows how to load from bytes in memory
      # return np.load(io.BytesIO(data))
      return np.load(data)
  
  
  def read_json(path):
      # Read and parse a JSON file from the given file path
      with open(path, 'r', encoding='utf-8') as file:
          data = json.load(file)
      return data
  
  
  def padding(image, mask, center=True, padding_ratio_range=[1.15, 1.15]):
      """
      Pad the input image and mask to a square shape with padding ratio.
  
      Args:
          image (np.ndarray): Input image array of shape (H, W, C).
          mask (np.ndarray): Corresponding mask array of shape (H, W).
          center (bool): Whether to center the original image in the padded output.
          padding_ratio_range (list): Range [min, max] to randomly select padding ratio.
  
      Returns:
          newimg (np.ndarray): Padded image of shape (resize_side, resize_side, 3).
          newmask (np.ndarray): Padded mask of shape (resize_side, resize_side).
      """
      h, w = image.shape[:2]
      max_side = max(h, w)
  
      # Select padding ratio either fixed or randomly within the given range
      if padding_ratio_range[0] == padding_ratio_range[1]:
          padding_ratio = padding_ratio_range[0]
      else:
          padding_ratio = random.uniform(padding_ratio_range[0], padding_ratio_range[1])
      resize_side = int(max_side * padding_ratio)
      # resize_side = int(max_side * 1.15)
  
      pad_h = resize_side - h
      pad_w = resize_side - w
      if center:
          start_h = pad_h // 2
      else:
          start_h = pad_h - resize_side // 20
          
      start_w = pad_w // 2
  
      # Create new white image and black mask with padded size
      newimg = np.ones((resize_side, resize_side, 3), dtype=np.uint8) * 255
      newmask = np.zeros((resize_side, resize_side), dtype=np.uint8)
      
      # Place original image and mask into the padded canvas
      newimg[start_h:start_h + h, start_w:start_w + w] = image
      newmask[start_h:start_h + h, start_w:start_w + w] = mask
      
      return newimg, newmask
  
  
  def viz_pc(surface, normal, image_input, name):
      image_input = image_input.cpu().numpy()
      image_input = image_input.transpose(1, 2, 0) * 0.5 + 0.5
      image_input = (image_input * 255).astype(np.uint8)
      cv2.imwrite(name + '.png', cv2.cvtColor(image_input, cv2.COLOR_RGB2BGR))
      surface = surface.cpu().numpy()
      normal = normal.cpu().numpy()
      surface_mesh = trimesh.Trimesh(surface, vertex_colors=(normal + 1) / 2)
      surface_mesh.export(name + '.obj')
  
  
  class AlignedShapeLatentDataset(torch.utils.data.dataset.IterableDataset):
      def __init__(
          self,
          data_list: str = None,
          cond_stage_key: str = "image",
          image_transform = None,
          pc_size: int = 2048,
          pc_sharpedge_size: int = 2048,
          sharpedge_label: bool = False,
          return_normal: bool = False,
          deterministic = False,
          worker_seed = None,
          padding = True,
          padding_ratio_range=[1.15, 1.15]
      ):
          super().__init__()
          if isinstance(data_list, str) and data_list.endswith('.json'):
              self.data_list = read_json(data_list_json)
          elif isinstance(data_list, str) and os.path.isdir(data_list):
              self.data_list = glob.glob(data_list + '/*')
          else:
              self.data_list = data_list
          assert isinstance(self.data_list, list)
          self.rng = random.Random(0)
          
          self.cond_stage_key = cond_stage_key
          self.image_transform = image_transform
          
          self.pc_size = pc_size
          self.pc_sharpedge_size = pc_sharpedge_size
          self.sharpedge_label = sharpedge_label
          self.return_normal = return_normal
  
          self.padding = padding
          self.padding_ratio_range = padding_ratio_range
          
          rank_zero_info(f'*' * 50)
          rank_zero_info(f'Dataset Infos:')
          rank_zero_info(f'# of 3D file: {len(self.data_list)}')
          rank_zero_info(f'# of Surface Points: {self.pc_size}')
          rank_zero_info(f'# of Sharpedge Surface Points: {self.pc_sharpedge_size}')
          rank_zero_info(f'Using sharp edge label: {self.sharpedge_label}')
          rank_zero_info(f'*' * 50)
  
  
      def load_surface_sdf_points(self, rng, random_surface, sharpedge_surface):
          surface_normal = []
          if self.pc_size > 0:
              ind = rng.choice(random_surface.shape[0], self.pc_size, replace=False)
              random_surface = random_surface[ind]
              if self.sharpedge_label:
                  sharpedge_label = np.zeros((self.pc_size, 1))
                  random_surface = np.concatenate((random_surface, sharpedge_label), axis=1)
              surface_normal.append(random_surface)
              
          if self.pc_sharpedge_size > 0:
              ind_sharpedge = rng.choice(sharpedge_surface.shape[0], self.pc_sharpedge_size, replace=False)
              sharpedge_surface = sharpedge_surface[ind_sharpedge]
              if self.sharpedge_label:
                  sharpedge_label = np.ones((self.pc_sharpedge_size, 1))
                  sharpedge_surface = np.concatenate((sharpedge_surface, sharpedge_label), axis=1)
              surface_normal.append(sharpedge_surface)
              
          surface_normal = np.concatenate(surface_normal, axis=0)
          surface_normal = torch.FloatTensor(surface_normal)
          surface = surface_normal[:, 0:3]
          normal = surface_normal[:, 3:6]
          assert surface.shape[0] == self.pc_size + self.pc_sharpedge_size
          
          geo_points = 0.0
          normal = torch.nn.functional.normalize(normal, p=2, dim=1)
          if self.return_normal:
              surface = torch.cat([surface, normal], dim=-1)
          if self.sharpedge_label:
              surface = torch.cat([surface, surface_normal[:, -1:]], dim=-1)
          return surface, geo_points
  
      def load_render(self, imgs_path):
          imgs_choice = self.rng.sample(imgs_path, 1)
          images, masks = [], []
          for image_path in imgs_choice:
              image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
              assert image.shape[2] == 4
              alpha = image[:, :, 3:4].astype(np.float32) / 255
              forground = image[:, :, :3]
              background = np.ones_like(forground) * 255
              img_new = forground * alpha + background * (1 - alpha)
              image = img_new.astype(np.uint8)
              image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
              mask = (alpha[:, :, 0] * 255).astype(np.uint8)
  
              if self.padding:
                  h, w = image.shape[:2]
                  binary = mask > 0.3
                  non_zero_coords = np.argwhere(binary)
                  x_min, y_min = non_zero_coords.min(axis=0)
                  x_max, y_max = non_zero_coords.max(axis=0)
                  image, mask = padding(
                      image[max(x_min - 5, 0):min(x_max + 5, h), max(y_min - 5, 0):min(y_max + 5, w)],
                      mask[max(x_min - 5, 0):min(x_max + 5, h), max(y_min - 5, 0):min(y_max + 5, w)],
                      center=True, padding_ratio_range=self.padding_ratio_range)
              
              if self.image_transform:
                  image = self.image_transform(image)
                  mask = np.stack((mask, mask, mask), axis=-1)
                  mask = self.image_transform(mask)
                  
              images.append(image)
              masks.append(mask)
              
          images = torch.cat(images, dim=0)
          masks = torch.cat(masks, dim=0)[:1, ...]
          return images, masks
  
      def decode(self, item):
          uid = item.split('/')[-1]
          render_img_paths = [os.path.join(item, f'render_cond/{i:03d}.png') for i in range(24)]
          # transforms_json_path = os.path.join(item, 'render_cond/transforms.json')
          surface_npz_path = os.path.join(item, f'geo_data/{uid}_surface.npz')
          # sdf_npz_path = os.path.join(item, f'geo_data/{uid}_sdf.npz')
          # watertight_obj_path = os.path.join(item, f'geo_data/{uid}_watertight.obj')
          sample = {}
          sample["image"] = render_img_paths
          surface_data = read_npz(surface_npz_path)
          sample["random_surface"] = surface_data['random_surface']
          sample["sharpedge_surface"] = surface_data['sharp_surface']
          return sample
  
      def transform(self, sample):
          rng = np.random.default_rng()
          random_surface = sample.get("random_surface", 0)
          sharpedge_surface = sample.get("sharpedge_surface", 0)
          image_input, mask_input = self.load_render(sample['image'])
          surface, geo_points = self.load_surface_sdf_points(rng, random_surface, sharpedge_surface)
          sample = {
              "surface": surface,
              "geo_points": geo_points,
              "image": image_input,
              "mask": mask_input,
          }
          return sample
  
      def __iter__(self):
          total_num = 0
          failed_num = 0
          for data in ResampledShards(self.data_list):
              total_num += 1
              if total_num % 1000 == 0:
                  print(f"Current failure rate of data loading:")
                  print(f"{failed_num}/{total_num}={failed_num/total_num}")
              try:
                  sample = self.decode(data)
                  sample = self.transform(sample)
              except Exception as err:
                  print(err)
                  failed_num += 1
                  continue
              yield sample
  
  
  class AlignedShapeLatentModule(LightningDataModule):
      def __init__(
          self,
          batch_size: int = 1,
          num_workers: int = 4,
          val_num_workers: int = 2,
          train_data_list: str = None,
          val_data_list: str = None,
          cond_stage_key: str = "all",
          image_size: int = 224,
          mean: Union[List[float], Tuple[float]] = (0.485, 0.456, 0.406),
          std: Union[List[float], Tuple[float]] = (0.229, 0.224, 0.225),
          pc_size: int = 2048,
          pc_sharpedge_size: int = 2048,
          sharpedge_label: bool = False,
          return_normal: bool = False, 
          padding = True,
          padding_ratio_range=[1.15, 1.15]
      ):
  
          super().__init__()
          self.batch_size = batch_size
          self.num_workers = num_workers
          self.val_num_workers = val_num_workers
  
          self.train_data_list = train_data_list
          self.val_data_list = val_data_list
          
          self.cond_stage_key = cond_stage_key
          self.image_size = image_size
          self.mean = mean
          self.std = std
          self.train_image_transform = transforms.Compose([
              transforms.ToTensor(),
              transforms.Resize(self.image_size),
              transforms.Normalize(mean=self.mean, std=self.std)])
          self.val_image_transform = transforms.Compose([
              transforms.ToTensor(),
              transforms.Resize(self.image_size),
              transforms.Normalize(mean=self.mean, std=self.std)])
  
          self.pc_size = pc_size
          self.pc_sharpedge_size = pc_sharpedge_size
          self.sharpedge_label = sharpedge_label
          self.return_normal = return_normal
  
          self.padding = padding
          self.padding_ratio_range = padding_ratio_range
          
      def train_dataloader(self):
          asl_params = {
              "data_list": self.train_data_list,
              "cond_stage_key": self.cond_stage_key,
              "image_transform": self.train_image_transform,
              "pc_size": self.pc_size,
              "pc_sharpedge_size": self.pc_sharpedge_size,
              "sharpedge_label": self.sharpedge_label,
              "return_normal": self.return_normal,
              "padding": self.padding,
              "padding_ratio_range": self.padding_ratio_range
          }
          dataset = AlignedShapeLatentDataset(**asl_params)
          return torch.utils.data.DataLoader(
              dataset,
              batch_size=self.batch_size,
              num_workers=self.num_workers,
              pin_memory=True,
              drop_last=True,
              worker_init_fn=worker_init_fn,
          )
  
      def val_dataloader(self):
          asl_params = {
              "data_list": self.val_data_list,
              "cond_stage_key": self.cond_stage_key,
              "image_transform": self.val_image_transform,
              "pc_size": self.pc_size,
              "pc_sharpedge_size": self.pc_sharpedge_size,
              "sharpedge_label": self.sharpedge_label,
              "return_normal": self.return_normal, 
              "padding": self.padding,
              "padding_ratio_range": self.padding_ratio_range
          }
          dataset = AlignedShapeLatentDataset(**asl_params)
          return torch.utils.data.DataLoader(
              dataset,
              batch_size=self.batch_size,
              num_workers=self.val_num_workers,
              pin_memory=True,
              drop_last=True,
              worker_init_fn=worker_init_fn,
          )
  