#try this with pytorch directly
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
from PIL import Image as PILImage
import pandas as pd
import os

# Define the path to the train thumbnails
trainPath = Path('/kaggle/input/UBC-OCEAN/train_images')

# Define the transform to apply to each image
transform = transforms.Compose([
    transforms.Resize(460),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(*imagenet_stats)
])


# Define a function to load an image from a file
def load_image(filename):
    image = PILImage.open(filename)
    image = transform(image)
    return image

# Define a function to load the train dataset
def load_train_dataset():
    # Load the train CSV file
    df = pd.read_csv("/kaggle/input/UBC-OCEAN/train.csv")
    # Create a list of tuples containing the image path and label
    image_list = [(os.path.join(trainPath, str(row['image_id']) + '.png'), row['label']) for index, row in df.iterrows()]
    # Create a dataset from the list of image tuples
    dataset = [(load_image(f), l) for f, l in image_list]
    return dataset

# Load the train dataset
train_dataset = load_train_dataset()

# Create a dataloader for the train dataset
train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)

# Define the model architecture
model = models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)

# Move the model to the GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Train the model
model.train()
for epoch in range(5):
    running_loss = 0.0
    for i, batch in enumerate(train_dataloader):
        # Get the inputs and labels
        inputs, labels = batch
        inputs, labels = inputs.to(device), labels.to(device)

        # Zero the parameter gradients
        optimizer.zero_grad()

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

        # Backward pass and optimization
        loss.backward()
        optimizer.step()

        # Print statistics
        running_loss += loss.item()
        if i % 100 == 99:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0

print('Finished training')

# Save the model
torch.save(model.state_dict(), 'ovarianCancerModel.pth')