Class BinaryFaceForensic

datasets/face_forensic_binary.py:16–77

@DATASETS.register_module()
class BinaryFaceForensic(CommonDataset):
    def __init__(self,
                 config,
                 split,
                 **kwargs):
        """
        @params:
        config: Dataset config
        split: train/val/test which directs to the split folders
        """
        self.split = split
        if kwargs is not None:
            for k,v in kwargs.items():
                if v is None:
                    raise ValueError(f'{k}:{v} retrieve a None value!')
                self.__setattr__(k, v)
        super().__init__(config, **kwargs)

        #Load data
        self.image_paths, self.labels, self.mask_paths = self._load_data(split)
        
        #Calling transform methods for inputs
        self.geo_transform = build_pipeline(config.TRANSFORM.geometry, PIPELINES)
        self.colorjitter_transform = build_pipeline(config.TRANSFORM.color, PIPELINES)

    def _load_data(self, split):
        from_file = self._cfg.FROM_FILE
        
        if not from_file:
            image_paths, labels, mask_paths = self._load_from_path(split)
        else:
            image_paths, labels, mask_paths = self._load_from_file(split)
        
        assert len(image_paths) != 0, "Image paths have not been loaded! Please check image directory!"
        assert len(labels) != 0, "Labels have not been loaded! Please check image folder path!"
        return image_paths, labels, mask_paths

    def _load_img(self, img_path):
        img = Image.open(img_path)
        img = np.array(img)
        return img

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        label = np.expand_dims(self.labels[idx], axis=-1)
        img = self._load_img(img_path)

        #Applying geo transform to inputs
        geo_transfomed = self.geo_transform(img)
        img_trans = geo_transfomed['image']
        
        #Applying color transform to inputs
        color_transfomed = self.colorjitter_transform(img_trans)
        img_trans = color_transfomed['image']