lvyilin / pytorch-fgvc-dataset Goto Github PK

    # Define paths to the MAT files
    if self.train:
        self.annotations_mat = os.path.join(root, 'cars_train_annos.mat')
        self.images_dir = os.path.join(root, 'cars_train')
    else:
        self.annotations_mat = os.path.join(root, 'cars_test_annos_withlabels.mat')
        self.images_dir = os.path.join(root, 'cars_test')

    # Load annotations
    self.annotations = self._load_annotations()

    if download:
        self.download()

def _load_annotations(self):
    # Load .mat file
    annos = scipy.io.loadmat(self.annotations_mat)
    # This might require custom parsing depending on the structure of your .mat file
    annotations = []
    for anno in annos['annotations'][0]:
        annotations.append({
            'fname': anno[-1][0],
            #'bbox': anno[0:4],
            'class': anno[4][0][0]
        })
    return annotations

def __getitem__(self, index):
    # Load image
    img_path = os.path.join(self.images_dir, self.annotations[index]['fname'])
    image = Image.open(img_path).convert('RGB')

    # Get bounding box
    #bbox = self.annotations[index]['bbox']
    # Here, you might want to crop or otherwise use the bounding box.

    # Get label
    target = self.annotations[index]['class']

    if self.transform is not None:
        image = self.transform(image)

    if self.target_transform is not None:
        target = self.target_transform(target)

    return image, target

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

def download(self):
    # Implement this method if you want to automatically download data
    pass