# Image transformations (VERY IMPORTANT for ResNet)
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # ResNet needs this
    transforms.ToTensor()
])

# Load dataset
dataset = datasets.ImageFolder(
    root='/content/drive/MyDrive/TrashNet',
    transform=transform
)

# Create DataLoader
train_loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=32,
    shuffle=True
)

# Number of classes
NUM_CLASSES = len(dataset.classes)

print("Classes:", dataset.classes)

# Load pretrained ResNet
model = models.resnet18(pretrained=True)

# Freeze all layers (optional but recommended)
for param in model.parameters():
    param.requires_grad = False

# Replace final layer
model.fc = nn.Linear(model.fc.in_features, NUM_CLASSES)

# Move to device
model = model.to(device)

print(model)

criterion = nn.CrossEntropyLoss()

# Only train last layer
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

EPOCHS = 5

for epoch in range(EPOCHS):
    model.train()
    running_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch [{epoch+1}/{EPOCHS}], Loss: {running_loss:.4f}")