

import pandas as pd
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
from torchvision.transforms import functional as TF
import os
train = pd.read_csv('train.csv')
test = pd.read_csv('test.csv')
train.head()
train['label'].unique()
map = {"queen":0, "rook":1,"bishop":2,'knight':3,"pawn":4}
rev_map = {}
for key,value in map.items():
    rev_map[value]=key
rev_map
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(os.path.join("./images",f"{train['image_path'][index]}"))
        if self.transform:
            img = self.transform(img)
        label = map[row['label']]
        return img,label
transform = transforms.Compose([
    transforms.Resize((256,256)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
train_ds=TrainDataset("./train.csv", transform=transform)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=0)
device = torch.device('cuda')
weights = models.ResNet50_Weights.IMAGENET1K_V2
model = models.resnet50(weights=weights)
model.fc = nn.Linear(2048, 5)
for param in model.parameters():
    param.requires_grad = True
for param in model.fc.parameters():
    param.requires_grad = True
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
for epoch in range(10):
    total_loss = 0
    for img,label in train_loader:
        img = img.to(device)
        label = label.to(device)
        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"epoch{epoch+1}--loss{total_loss/len(train_loader)}")
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(os.path.join("./images",f"{test['image_path'][index]}"))
        if self.transform:
            img = self.transform(img)
        return img
test_ds=TestDataset("./test.csv", transform=transform)
test_loader = DataLoader(test_ds, batch_size=32, shuffle=False, num_workers=0)
preds = []
for img in test_loader:
    with torch.no_grad():
        img= img.to(device)
        pred = model(img)
        pred = torch.argmax(pred,1)
        preds.append(pred.cpu())
        preds_final = torch.cat(preds)
preds_final = preds_final.numpy()
preds_final
pred_final_rev = []
for i in preds_final:
    pred_final_rev.append(rev_map[i])
subs = pd.DataFrame({
    'id': test['id'],
    'label':pred_final_rev
}).to_csv('subs.csv',index=False)
