Update tasks/image.py

tasks/image.py

@@ -128,23 +128,23 @@ async def evaluate_image(request: ImageEvaluationRequest):
             x = self.fc2(x)
             return x
 
-    class CustomDataset(Dataset, labels):
-        def __init__(self, dataset, transform=None):
-            self.dataset = dataset
-            self.transform = transform
-            self.labels = labels
+    # class CustomDataset(Dataset, labels):
+    #     def __init__(self, dataset, transform=None):
+    #         self.dataset = dataset
+    #         self.transform = transform
+    #         self.labels = labels
 
-        def __len__(self):
-            return len(self.dataset)
+    #     def __len__(self):
+    #         return len(self.dataset)
 
-        def __getitem__(self, idx):
-            image = self.dataset[idx]['image']
-            label = self.labels[idx] 
+    #     def __getitem__(self, idx):
+    #         image = self.dataset[idx]['image']
+    #         label = self.labels[idx] 
 
-            if self.transform:
-                image = self.transform(image)
+    #         if self.transform:
+    #             image = self.transform(image)
 
-            return image, label
+    #         return image, label
 
     # Create an instance of the model
     model = ImageClassifier()