content: # 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 time
  import glob
  import json
  import random
  import numpy as np
  import torch
  from .loader_util import BaseDataset
  
  
  class TextureDataset(BaseDataset):
  
      def __init__(
          self, json_path, num_view=6, image_size=512, lighting_suffix_pool=["light_PL", "light_AL", "light_ENVMAP"]
      ):
          self.data = list()
          self.num_view = num_view
          self.image_size = image_size
          self.lighting_suffix_pool = lighting_suffix_pool
          if isinstance(json_path, str):
              json_path = [json_path]
          for jp in json_path:
              with open(jp) as f:
                  self.data.extend(json.load(f))
          print("============= length of dataset %d =============" % len(self.data))
  
      def __getitem__(self, index):
          try_sleep_interval = 20
          total_try_num = 100
          cnt = try_sleep_interval * total_try_num
          # try:
          images_ref = list()
          images_albedo = list()
          images_mr = list()
          images_normal = list()
          images_position = list()
          bg_white = [1.0, 1.0, 1.0]
          bg_black = [0.0, 0.0, 0.0]
          bg_gray = [127 / 255.0, 127 / 255.0, 127 / 255.0]
          dirx = self.data[index]
  
          condition_dict = {}
  
          # 6view
          fix_num_view = self.num_view
          available_views = []
          for ext in ["*_albedo.png", "*_albedo.jpg", "*_albedo.jpeg"]:
              available_views.extend(glob.glob(os.path.join(dirx, "render_tex", ext)))
          cond_images = (
              glob.glob(os.path.join(dirx, "render_cond", "*.png"))
              + glob.glob(os.path.join(dirx, "render_cond", "*.jpg"))
              + glob.glob(os.path.join(dirx, "render_cond", "*.jpeg"))
          )
  
          # 确保有足够的样本
          if len(available_views) < fix_num_view:
              print(
                  f"Warning: Only {len(available_views)} views available, but {fix_num_view} requested."
                  "Using all available views."
              )
              images_gen = available_views
          else:
              images_gen = random.sample(available_views, fix_num_view)
  
          if not cond_images:
              raise ValueError(f"No condition images found in {os.path.join(dirx, 'render_cond')}")
          ref_image_path = random.choice(cond_images)
          light_suffix = None
          for suffix in self.lighting_suffix_pool:
              if suffix in ref_image_path:
                  light_suffix = suffix
                  break
          if light_suffix is None:
              raise ValueError(f"light suffix not found in {ref_image_path}")
          ref_image_diff_light_path = random.choice(
              [
                  ref_image_path.replace(light_suffix, tar_suffix)
                  for tar_suffix in self.lighting_suffix_pool
                  if tar_suffix != light_suffix
              ]
          )
          images_ref_paths = [ref_image_path, ref_image_diff_light_path]
  
          # Data aug
          bg_c_record = None
          for i, image_ref in enumerate(images_ref_paths):
              if random.random() < 0.6:
                  bg_c = bg_gray
              else:
                  if random.random() < 0.5:
                      bg_c = bg_black
                  else:
                      bg_c = bg_white
              if i == 0:
                  bg_c_record = bg_c
              image, alpha = self.load_image(image_ref, bg_c_record)
              image = self.augment_image(image, bg_c_record).float()
              images_ref.append(image)
          condition_dict["images_cond"] = torch.stack(images_ref, dim=0).float()
  
          for i, image_gen in enumerate(images_gen):
              images_albedo.append(self.augment_image(self.load_image(image_gen, bg_gray)[0], bg_gray))
              images_mr.append(
                  self.augment_image(self.load_image(image_gen.replace("_albedo", "_mr"), bg_gray)[0], bg_gray)
              )
              images_normal.append(
                  self.augment_image(self.load_image(image_gen.replace("_albedo", "_normal"), bg_gray)[0], bg_gray)
              )
              images_position.append(
                  self.augment_image(self.load_image(image_gen.replace("_albedo", "_pos"), bg_gray)[0], bg_gray)
              )
  
          condition_dict["images_albedo"] = torch.stack(images_albedo, dim=0).float()
          condition_dict["images_mr"] = torch.stack(images_mr, dim=0).float()
          condition_dict["images_normal"] = torch.stack(images_normal, dim=0).float()
          condition_dict["images_position"] = torch.stack(images_position, dim=0).float()
          condition_dict["name"] = dirx  # .replace('/', '_')
          return condition_dict  # (N, 3, H, W)
  
          # except Exception as e:
          #     print(e, self.data[index])
          #     # exit()
  
  
  if __name__ == "__main__":
      dataset = TextureDataset(json_path=["../../../train_examples/examples.json"])
      print("images_cond", dataset[0]["images_cond"].shape)
      print("images_albedo", dataset[0]["images_albedo"].shape)
      print("images_mr", dataset[0]["images_mr"].shape)
      print("images_normal", dataset[0]["images_normal"].shape)
      print("images_position", dataset[0]["images_position"].shape)
      print("name", dataset[0]["name"])
  