

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms
from PIL import Image
import pandas as pd
train_df = pd.read_csv('train.csv')
train_df.head()
class TrainDataset(Dataset):
    def __init__(self,path,transform):
        self.df = pd.read_csv(path)
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        img = Image.open(row['image_path']).convert("RGB")
        label = row['label']
        if self.transform:
            img = self.transform(img)
        return img,label
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
])
train_ds = TrainDataset("./train.csv", transform=transform)
train = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=0)
train_df['label'].unique()
device = torch.device('cuda')
weights = models.ResNet50_Weights.IMAGENET1K_V1
model = models.resnet50(weights=weights)
model.fc = nn.Linear(2048,2)
model.train()
for param in model.parameters():
    param.requires_grad = False
for param in model.layer4.parameters():
    param.requires_grad = True
for param in model.fc.parameters():
    param.requires_grad = True

optimizer = optim.Adam([
    {'params': model.layer4.parameters(), 'lr': 1e-4},
    {'params': model.fc.parameters(), 'lr': 1e-3}
])
model = model.to(device)
cirterion = nn.CrossEntropyLoss()
for epoch in range(20):
    for image, label in train:
        image = image.to(device)
        label = label.to(device)
        optimizer.zero_grad()
        output = model(image)
        loss= cirterion(output,label)
        loss.backward()
        optimizer.step()
    print(f"EPOCH{epoch+1} ------ LOSS{loss:04f}")
class TestDataset(Dataset):
    def __init__(self,path,transform):
        self.df = pd.read_csv(path)
        self.transform = transform
    def __len__(self):
        return len(self.df)
    def __getitem__(self, index):
        row = self.df.iloc[index]
        img = Image.open(row['image_path']).convert("RGB")
        # label = row['label']
        if self.transform:
            img = self.transform(img)
        return img
test_ds = TestDataset("./test.csv", transform=transform)
test= DataLoader(test_ds, batch_size=32, shuffle=True, num_workers=0)
model.eval()
preds=[]
preds_final = []
for img in test:
    img = img.to(device)
    with torch.no_grad():
        pred = model(img)
        _,pred = torch.max(pred, 1)
        preds.append(pred.cpu())
        preds_final = torch.cat(preds)
preds_final
len(preds_final)
preds_final = preds_final.numpy()
test_df = pd.read_csv('test.csv')
len(test_df)
subs = pd.DataFrame({
    'image_path': test_df['image_path'],
    'label':preds_final
}).to_csv("subs.csv",index=False)
