Car_VS_Rest / dataloder_pytorch.py
Nekshay's picture
Update dataloder_pytorch.py
cf4f71b verified
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.