| import torch |
| from torch import nn |
| from torch.utils.data import DataLoader |
|
|
| # Hyperparameters |
| image_size = (224, 224, 3) # Adjust based on your data |
|
|
| # Define the Generator Network |
| class Generator(nn.Module): |
| def __init__(self): |
| super(Generator, self).__init__() |
| # Define convolutional layers with appropriate filters and activations |
| self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1) |
| # ... Add more convolutional layers as needed |
| self.conv_final = nn.Conv2d(128, 3, kernel_size=3, stride=1, padding=1, activation=nn.Tanh) # Tanh for shadow intensity |
|
|
| def forward(self, x): |
| # Define the forward pass through the convolutional layers |
| x = self.conv1(x) |
| # ... Forward pass through remaining convolutional layers |
| return self.conv_final(x) |
|
|
| # Define the Discriminator Network |
| class Discriminator(nn.Module): |
| def __init__(self): |
| super(Discriminator, self).__init__() |
| # Define convolutional layers with appropriate filters and activations |
| self.conv1 = nn.Conv2d(6, 32, kernel_size=3, stride=1, padding=1) |
| # ... Add more convolutional layers as needed |
| self.linear = nn.Linear(128, 1) # Final layer with sigmoid activation |
|
|
| def forward(self, car, shadow): |
| # Concatenate car and shadow features |
| x = torch.cat([car, shadow], dim=1) |
| # Define the forward pass through the convolutional layers |
| x = self.conv1(x) |
| # ... Forward pass through remaining convolutional layers |
| return torch.sigmoid(self.linear(x)) |
|
|
| # Create data loaders for training and validation data |
| # ... (Implement data loading logic using PyTorch's DataLoader) |
|
|
| # Create the models |
| generator = Generator() |
| discriminator = Discriminator() |
|
|
| # Define loss function and optimizer |
| criterion = nn.BCELoss() |
| g_optimizer = torch.optim.Adam(generator.parameters(), lr=0.0002) |
| d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=0.0002) |
|
|
| # Training loop |
| for epoch in range(epochs): |
| # Train the Discriminator |
| # ... (Implement discriminator training logic with loss calculation and updates) |
|
|
| # Train the Generator |
| # ... (Implement generator training logic with loss calculation and updates) |
|
|
| # Print training progress |
| # ... (Print loss values or other metrics) |
|
|
| # Save the trained generator |
| torch.save(generator.state_dict(), 'generator.pt') |
|
|
