Update dataloder_pytorch.py

dataloder_pytorch.py

@@ -0,0 +1,48 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+# Define a custom Dataset class
+class CarShadowDataset(Dataset):
+  def __init__(self, root_dir, transform=None):
+    self.root_dir = root_dir
+    self.transform = transform
+    self.image_paths = []  # List to store image paths
+
+    # Loop through car and shadow folders to collect image paths
+    for phase in ['train', 'val', 'test']:  # Adjust based on your data structure
+      car_folder = os.path.join(root_dir, phase, 'car')
+      shadow_folder = os.path.join(root_dir, phase, 'shadow')
+
+      for filename in os.listdir(car_folder):
+        car_path = os.path.join(car_folder, filename)
+        shadow_path = os.path.join(shadow_folder, filename.split('.')[0] + '_shadow.jpg')  # Assuming consistent naming
+        self.image_paths.append((car_path, shadow_path))
+
+  def __len__(self):
+    return len(self.image_paths)
+
+  def __getitem__(self, idx):
+    car_path, shadow_path = self.image_paths[idx]
+    car_image = load_image(car_path)  # Replace with your image loading function
+    shadow_image = load_image(shadow_path)  # Replace with your image loading function
+
+    if self.transform:
+      car_image = self.transform(car_image)
+      shadow_image = self.transform(shadow_image)
+
+    return car_image, shadow_image
+
+# Function to load image (replace with your preferred method)
+def load_image(path):
+  # Implement image loading using libraries like OpenCV or PIL
+  # Ensure images are converted to tensors and normalized if needed
+  # ...
+
+# Prepare data loaders
+train_data = DataLoader(CarShadowDataset(root_dir='dataset/train', transform=your_transform), batch_size=32, shuffle=True)
+val_data = DataLoader(CarShadowDataset(root_dir='dataset/val', transform=your_transform), batch_size=32)  # Optional for validation
+
+# Example usage
+for car_image, shadow_image in train_data:
+  # Access your data for training
+  # ...