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
  # ...


# ... (Previous code for model definition and DataLoader)

# Discriminator Training
def train_discriminator(d_optimizer, real_images, fake_images, real_labels, fake_labels):
  # Clear gradients
  d_optimizer.zero_grad()

  # Forward pass through discriminator
  d_real_output = discriminator(real_images, real_images)  # Real images with real shadows
  d_fake_output = discriminator(real_images, fake_images)  # Real images with generated shadows

  # Calculate loss
  d_real_loss = criterion(d_real_output, torch.ones_like(d_real_output))
  d_fake_loss = criterion(d_fake_output, torch.zeros_like(d_fake_output))
  d_loss = (d_real_loss + d_fake_loss) / 2

  # Backpropagate and update weights
  d_loss.backward()
  d_optimizer.step()

  # Return the discriminator loss
  return d_loss.item()

# Generator Training
def train_generator(g_optimizer, real_images, fake_images):
  # Clear gradients
  g_optimizer.zero_grad()

  # Forward pass through discriminator (using generated shadows)
  g_fake_output = discriminator(real_images, fake_images)

  # Calculate loss (try to fool the discriminator)
  g_loss = criterion(g_fake_output, torch.ones_like(g_fake_output))

  # Backpropagate and update weights
  g_loss.backward()
  g_optimizer.step()

  # Return the generator loss
  return g_loss.item()

# Training loop
for epoch in range(epochs):
  for i, (real_images, real_shadows) in enumerate(train_data):
    # Generate fake shadows
    fake_shadows = generator(real_images)

    # Train discriminator
    d_loss = train_discriminator(d_optimizer, real_images, fake_shadows, torch.ones(real_images.size(0)), torch.zeros(real_images.size(0)))

    # Train generator
    g_loss = train_generator(g_optimizer, real_images, fake_shadows)

    # Print training progress
    if i % 100 == 0:
      print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_data)}], D_loss: {d_loss:.4f}, G_loss: {g_loss:.4f}')

I'd like to commend my colleague, Debashish, for his outstanding performance and invaluable contributions to our team. Debashish's deep understanding and expertise in machine learning concepts have been instrumental in driving our projects forward. He is always enthusiastic about experimenting with new ideas and technologies, which has led to innovative solutions and advancements in our work. Notably, Debashish developed a highly efficient tiny image classifier tailored for mobile applications and successfully integrated this ML model into our mobile app. Additionally, his efforts were crucial in creating a proof of concept (POC) for an ML-integrated mobile app using React-Native. Debashish's dedication and technical prowess have significantly enhanced our capabilities and project outcomes.