def __init__()

in models/datasets/attrib_dataset.py [0:0]

• 44 lines of code
• 14 McCabe index (conditional complexity)

    def __init__(self,
                 pathdb,
                 attribDictPath=None,
                 specificAttrib=None,
                 transform=None,
                 mimicImageFolder=False,
                 ignoreAttribs=False,
                 getEqualizer=False,
                 pathMask=None):
        r"""
        Args:

            - root (string): path to the directory containing the images
            - attribDictPath (string): path to a json file containing the images'
                                       specific attributes
            - specificAttrib (list of string): if not None, specify which attributes
                                                be selected
            - transform (torchvision.transforms): transformation to apply to the
                                                  loaded images.
            - mimicImageFolder (bool): set to True if the dataset is stored in the
                                      torchvision.datasets.ImageFolder format
            - ignoreAttribs (bool): set to True if you just want to use the attrib
                                    dict as a filter on images' name
        """

        self.totAttribSize = 0
        self.hasAttrib = attribDictPath is not None or mimicImageFolder
        self.pathdb = pathdb
        self.transform = transform
        self.shiftAttrib = None
        self.stats = None
        self.pathMask = None

        if attribDictPath:
            if ignoreAttribs:
                self.attribDict = None

                with open(attribDictPath, 'rb') as file:
                    tmpDict = json.load(file)
                self.listImg = [imgName for imgName in os.listdir(pathdb)
                                if (os.path.splitext(imgName)[1] in [".jpg",
                                                                     ".png", ".npy"] and imgName in tmpDict)]
            else:
                self.loadAttribDict(attribDictPath, pathdb, specificAttrib)

        elif mimicImageFolder:
            self.loadImageFolder(pathdb)
        else:
            self.attribDict = None
            self.listImg = [imgName for imgName in os.listdir(pathdb)
                            if os.path.splitext(imgName)[1] in [".jpg", ".png",
                                                                ".npy"]]

        if pathMask is not None:
            print("Path mask found " + pathMask)
            self.pathMask = pathMask
            self.listImg = [imgName for imgName in self.listImg
                            if os.path.isfile(os.path.join(pathMask,
                                                           os.path.splitext(imgName)[0] + "_mask.jpg"))]

        if len(self.listImg) == 0:
            raise AttributeError("Empty dataset")

        self.buildStatsOnDict()

        print("%d images found" % len(self))