class DatasetPersonDetection(DatasetPerson):
    """Dataset containing images from only one person with face detection"""

    def __init__(self, root_dir, transform=None):
        """
        :param root_dir: Directory with the images.
        :param transform: Transformations applied to the images.
        """
        # don't transform the images -> transform=None
        super().__init__(root_dir, transform=None)
        self.transform = transform
        self.detected_faces = []

        # run face detection on images
        for img in self.images:
            face_location = face_recognition.face_locations(img, model='hog')

            # ignore if 2 faces detected because in most cases they originate not form the same person
            if face_location and len(face_location) == 1:
                top, right, bottom, left = face_location[0]
                cropped_img = img[top:bottom, left:right]

                if self.transform:
                    cropped_img = self.transform(cropped_img)

                self.detected_faces.append(cropped_img)
        self.images.clear()
        self.images = self.detected_faces