(self,
difflocks_path,
processed_difflocks_path=None, #contains path to scalp textures
train=None,
load_rgb_imgs=False,
load_hair_mask=False,
load_orientation_map=False,
load_density_imgs=False,
load_full_strands=False,
load_guide_strands=False,
load_interpolated_strands=False,
load_cam=False,
load_path=True,
load_material=False,
load_scalp_texture=False,
scalp_texture_resolution = 256,
load_latents=False,
latents_type_list=["sdxl"],
load_latents_layers=[],
compute_tbn_full_strands=False,
compute_tbn_guide_strands=False,
compute_tbn_interpolated_strands=False,
# compute_fourier_full_strands=False,
nr_full_strands_per_hairstyle=None,
subsample_factor=None,
check_validity=True,
check_validity_processed=True,
do_pedantic_checks=True, #some checks may be redundant when doing on lambda where we know that numpy_state is the last thing to write, but may be worth doing locally where the data is not fully downloaded
overfit=False,
restrict_to_single_hairstyle_name=None,
train_ratio = 0.9,
randomly_flip=False,
max_nr_samples=None
)
| 122 | |
| 123 | class DiffLocksDataset(Dataset): |
| 124 | def __init__(self, |
| 125 | difflocks_path, |
| 126 | processed_difflocks_path=None, #contains path to scalp textures |
| 127 | train=None, |
| 128 | load_rgb_imgs=False, |
| 129 | load_hair_mask=False, |
| 130 | load_orientation_map=False, |
| 131 | load_density_imgs=False, |
| 132 | load_full_strands=False, |
| 133 | load_guide_strands=False, |
| 134 | load_interpolated_strands=False, |
| 135 | load_cam=False, |
| 136 | load_path=True, |
| 137 | load_material=False, |
| 138 | load_scalp_texture=False, |
| 139 | scalp_texture_resolution = 256, |
| 140 | load_latents=False, |
| 141 | latents_type_list=["sdxl"], |
| 142 | load_latents_layers=[], |
| 143 | compute_tbn_full_strands=False, |
| 144 | compute_tbn_guide_strands=False, |
| 145 | compute_tbn_interpolated_strands=False, |
| 146 | # compute_fourier_full_strands=False, |
| 147 | nr_full_strands_per_hairstyle=None, |
| 148 | subsample_factor=None, |
| 149 | check_validity=True, |
| 150 | check_validity_processed=True, |
| 151 | do_pedantic_checks=True, #some checks may be redundant when doing on lambda where we know that numpy_state is the last thing to write, but may be worth doing locally where the data is not fully downloaded |
| 152 | overfit=False, |
| 153 | restrict_to_single_hairstyle_name=None, |
| 154 | train_ratio = 0.9, |
| 155 | randomly_flip=False, |
| 156 | max_nr_samples=None |
| 157 | ): |
| 158 | |
| 159 | if subsample_factor is not None: |
| 160 | assert 768% (1.0/subsample_factor) == 0 #The subsample factor does not cleanly divide the image size |
| 161 | |
| 162 | self.difflocks_path=difflocks_path |
| 163 | self.processed_difflocks_path=processed_difflocks_path |
| 164 | self.train=train |
| 165 | self.load_rgb_imgs=load_rgb_imgs |
| 166 | self.load_hair_mask=load_hair_mask |
| 167 | self.load_orientation_map=load_orientation_map |
| 168 | self.load_density_imgs= load_density_imgs |
| 169 | self.load_full_strands=load_full_strands |
| 170 | self.load_guide_strands=load_guide_strands |
| 171 | self.load_interpolated_strands=load_interpolated_strands |
| 172 | self.load_cam=load_cam |
| 173 | self.load_path=load_path |
| 174 | self.load_material=load_material |
| 175 | self.load_scalp_texture = load_scalp_texture |
| 176 | self.scalp_texture_resolution=scalp_texture_resolution |
| 177 | self.load_latents=load_latents |
| 178 | self.load_latents_layers= load_latents_layers |
| 179 | self.latents_type_list=latents_type_list |
| 180 | self.compute_tbn_full_strands=compute_tbn_full_strands |
| 181 | self.compute_tbn_guide_strands=compute_tbn_guide_strands |
nothing calls this directly
no test coverage detected