Upload dataset.py with huggingface_hub

dataset.py

@@ -0,0 +1,52 @@
+import os
+from torchvision import transforms
+import pandas as pd
+import numpy as np
+from torch.utils.data import Dataset
+from torchvision.datasets.folder import default_loader
+
+IMAGE_NET_MEAN = [0.485, 0.456, 0.406]
+IMAGE_NET_STD = [0.229, 0.224, 0.225]
+normalize = transforms.Normalize(
+    mean=IMAGE_NET_MEAN,
+    std=IMAGE_NET_STD)
+
+class AVADataset(Dataset):
+    def __init__(self, path_to_csv, images_path, if_train):
+        self.df = pd.read_csv(path_to_csv)
+        self.images_path = images_path
+        self.if_train = if_train
+        if if_train:
+            self.transform = transforms.Compose([
+                # transforms.Resize((256, 256)),
+                transforms.RandomHorizontalFlip(),
+                # transforms.RandomCrop((224, 224)),
+                transforms.ToTensor(),
+                normalize])
+        else:
+            self.transform = transforms.Compose([
+                # transforms.Resize((224, 224)),
+                transforms.ToTensor(),
+                normalize])
+
+    def __len__(self):
+        return self.df.shape[0]
+
+    def __getitem__(self, item):
+        row = self.df.iloc[item]
+
+        scores_names = [f'score{i}' for i in range(2, 12)]
+        y = np.array([row[k] for k in scores_names])
+
+        p = y / y.sum()
+        image_id = row['image_id']
+
+        image_id = int(image_id)
+
+        image_path = os.path.join(self.images_path, f'{image_id}.jpg')
+        image = default_loader(image_path)
+
+        image = image.resize((224, 224))
+        x = self.transform(image)
+
+        return x, p.astype('float32')