Tanner Collin
c1d11ea3f7
Co-authored-by: aider (gemini/gemini-2.5-pro-preview-05-06) <aider@aider.chat>
113 lines
4.2 KiB
Python
113 lines
4.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader, random_split, WeightedRandomSampler
|
|
from torchvision import datasets, transforms
|
|
from PIL import Image, ImageDraw
|
|
import os
|
|
|
|
from model import (CropLowerRightTriangle, GarageDoorCNN, TRIANGLE_CROP_WIDTH,
|
|
TRIANGLE_CROP_HEIGHT, RESIZE_DIM)
|
|
|
|
|
|
def train_model():
|
|
# --- Hyperparameters and Configuration ---
|
|
DATA_DIR = 'data/labelled'
|
|
MODEL_SAVE_PATH = 'garage_door_cnn.pth'
|
|
NUM_EPOCHS = 10
|
|
BATCH_SIZE = 32
|
|
LEARNING_RATE = 0.001
|
|
|
|
# --- Data Preparation ---
|
|
# Define transforms
|
|
data_transforms = transforms.Compose([
|
|
CropLowerRightTriangle(triangle_width=TRIANGLE_CROP_WIDTH, triangle_height=TRIANGLE_CROP_HEIGHT),
|
|
transforms.Resize((RESIZE_DIM, RESIZE_DIM)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
|
])
|
|
|
|
# Load dataset with ImageFolder
|
|
full_dataset = datasets.ImageFolder(DATA_DIR, transform=data_transforms)
|
|
|
|
# Split into training and validation sets
|
|
train_size = int(0.8 * len(full_dataset))
|
|
val_size = len(full_dataset) - train_size
|
|
train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])
|
|
|
|
# --- Handle Class Imbalance ---
|
|
# Get labels for training set
|
|
train_labels = [full_dataset.targets[i] for i in train_dataset.indices]
|
|
|
|
# Get class counts
|
|
class_counts = torch.bincount(torch.tensor(train_labels))
|
|
|
|
# Compute weight for each class (inverse of count)
|
|
class_weights = 1. / class_counts.float()
|
|
|
|
# Assign a weight to each sample in the training set
|
|
sample_weights = torch.tensor([class_weights[label] for label in train_labels])
|
|
|
|
# Create a WeightedRandomSampler to balance the classes during training
|
|
sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights), replacement=True)
|
|
|
|
# The sampler will handle shuffling, so shuffle must be False for the DataLoader
|
|
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=sampler)
|
|
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
|
|
|
|
# --- Model, Loss, Optimizer ---
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Using device: {device}")
|
|
|
|
model = GarageDoorCNN(resize_dim=RESIZE_DIM).to(device)
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
|
|
|
|
# --- Training Loop ---
|
|
print("Starting training...")
|
|
for epoch in range(NUM_EPOCHS):
|
|
model.train()
|
|
running_loss = 0.0
|
|
for inputs, labels in train_loader:
|
|
inputs, labels = inputs.to(device), labels.to(device)
|
|
|
|
optimizer.zero_grad()
|
|
outputs = model(inputs)
|
|
loss = criterion(outputs, labels)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
running_loss += loss.item() * inputs.size(0)
|
|
|
|
epoch_loss = running_loss / len(train_dataset)
|
|
print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Training Loss: {epoch_loss:.4f}")
|
|
|
|
# --- Validation Loop ---
|
|
model.eval()
|
|
val_loss = 0.0
|
|
corrects = 0
|
|
with torch.no_grad():
|
|
for inputs, labels in val_loader:
|
|
inputs, labels = inputs.to(device), labels.to(device)
|
|
outputs = model(inputs)
|
|
loss = criterion(outputs, labels)
|
|
val_loss += loss.item() * inputs.size(0)
|
|
_, preds = torch.max(outputs, 1)
|
|
corrects += torch.sum(preds == labels.data)
|
|
|
|
val_epoch_loss = val_loss / len(val_dataset)
|
|
val_epoch_acc = corrects.double() / len(val_dataset)
|
|
print(f"Validation Loss: {val_epoch_loss:.4f}, Accuracy: {val_epoch_acc:.4f}")
|
|
|
|
# --- Save the trained model ---
|
|
torch.save(model.state_dict(), MODEL_SAVE_PATH)
|
|
print(f"Model saved to {MODEL_SAVE_PATH}")
|
|
|
|
if __name__ == '__main__':
|
|
# Check if data directory exists
|
|
if not os.path.isdir('data/labelled/open') or not os.path.isdir('data/labelled/closed'):
|
|
print("Error: Data directories 'data/open' and 'data/closed' not found.")
|
|
print("Please create them and place your image snapshots inside.")
|
|
else:
|
|
train_model()
|