Upload New File

4_Resnet_Solar_panel_defect_Identification/data.py

+from torch.utils.data import Dataset
+import torch
+from pathlib import Path
+from skimage.io import imread
+from skimage.color import gray2rgb
+import numpy as np
+import torchvision as tv
+
+train_mean = [0.59685254, 0.59685254, 0.59685254]
+train_std = [0.16043035, 0.16043035, 0.16043035]
+
+
+class ChallengeDataset(Dataset):
+    # TODO implement the Dataset class according to the description
+    def __init__(self, data, mode):
+        super().__init__()
+        self.data = data
+        self.mode = mode
+        self.length = None
+        if mode == "train":
+            self._transform = tv.transforms.Compose(
+                [tv.transforms.ToPILImage(), tv.transforms.RandomHorizontalFlip(p=0.25),
+                 tv.transforms.ToTensor(), tv.transforms.Normalize(train_mean, train_std)])
+            # self._transform = tv.transforms.Compose([tv.transforms.ToPILImage(), tv.transforms.RandomVerticalFlip(p=0.2), tv.transforms.RandomRotation(degrees=30), tv.transforms.GaussianBlur(kernel_size=3), tv.transforms.ToTensor(), tv.transforms.Normalize(train_mean, train_std)])
+        else:
+            self._transform = tv.transforms.Compose([tv.transforms.ToPILImage(), tv.transforms.ToTensor(), tv.transforms.Normalize(train_mean, train_std)])
+
+    def __len__(self):
+        self.length = self.data.shape[0]
+        return self.length
+
+    def __getitem__(self, index):
+        image = imread((self.data.iloc[index])["filename"])
+        rgb_image = gray2rgb(image)
+        image_tensor = self._transform(rgb_image).type(torch.Tensor)
+        label = (self.data.iloc[index])[["crack", "inactive"]]
+        label_np = label.to_numpy().astype(float)
+        label_tensor = torch.from_numpy(label_np).type(torch.Tensor)
+        return image_tensor, label_tensor
+