Method __getitem__

data_loader/dataloader.py:397–643

(self, idx)

        return len(self.hairstyles_path_list)

    def __getitem__(self, idx):
        out_dict = {}

        if self.overfit==True:
            # idx=4
            idx=5
            np.random.seed(0)

        hairstyle_data= self.hairstyles_path_data_list[idx]
        hairstyle_path = hairstyle_data.hairstyle_path

        #if we need any of the processed data, we load also the path for it
        # if self.load_scalp_texture:
        if self.processed_difflocks_path is not None:
            hairstyle_path_processed = os.path.join(self.processed_difflocks_path, "processed_hairstyles", os.path.basename(hairstyle_path))

        # print("hairstyle_path",hairstyle_path)
            

        flip = False
        if self.randomly_flip:
            if random.random()<0.5:
                flip = True

            
        #read metadata
        with open(os.path.join(hairstyle_path,"metadata.json"),'r') as f:
            try:
                meta_data = f.read()
                meta_data = json.loads(meta_data)
                if isinstance(meta_data, str):
                    print("why is this a str", hairstyle_path)
                    exit(1)
            except:
                print("couldn't load metadata for", hairstyle_path)
                exit(1)


        if self.load_rgb_imgs:
            img_path = os.path.join(hairstyle_path,"rgb.png")
            img = read_image(img_path).float()/255.0
            if flip:
                img=torchvision.transforms.functional.hflip(img)
            if self.subsample_factor is not None:
                scale_factor=[1.0/self.subsample_factor, 1.0/self.subsample_factor]
                img=torch.nn.functional.interpolate(img.unsqueeze(0), scale_factor=scale_factor, mode="area").squeeze(0)
            out_dict["rgb_img"] = img

        if self.load_hair_mask:
            img_path = os.path.join(hairstyle_path_processed,"hair_masks","hair_mask.png")
            if flip:
                img_path = os.path.join(hairstyle_path_processed,"hair_masks_flip","hair_mask.png")
            img = read_image(img_path).float()/255.0
            img = img[0:1,:,:]
            if self.subsample_factor is not None:
                scale_factor=[1.0/self.subsample_factor, 1.0/self.subsample_factor]
                img=torch.nn.functional.interpolate(img.unsqueeze(0), scale_factor=scale_factor, mode="nearest").squeeze(0)
            out_dict["hair_mask"] = img