

import pandas as pd
train_df = pd.read_csv("train.csv")
test_df = pd.read_csv("test.csv")
train_df.head()
train_df['ImagePath'][0].split("/")[1].strip(".png")
train_df['Type'].unique()
map = {
    'Psychic': 0, 'Water': 1, 'Ghost': 2, 'Fighting': 3, 'Ice': 4,
    'Grass': 5, 'Normal': 6, 'Rock': 7, 'Electric': 8, 'Poison': 9,
    'Bug': 10, 'Fairy': 11, 'Fire': 12, 'Dark': 13, 'Steel': 14,
    'Dragon': 15, 'Ground': 16, 'Flying': 17
}
map_rev= {}
for key,value in map.items():
    map_rev[value] = key
map_rev
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
class TrainDataset(Dataset):
    def __init__(self, cv_path, transform):
        self.df = pd.read_csv(cv_path)
        self.transform = transform
    def __len__(self):
        return (len(self.df))
    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img= Image.open(row['ImagePath']).convert('RGB')
        if self.transform:
            img = self.transform(img)
        label = map[row['Type']]
        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=8, shuffle=False, num_workers=0)
device = torch.device('cuda')
weights = models.ResNet50_Weights.IMAGENET1K_V1
model = models.resnet50(weights=weights).eval()
model.fc = nn.Linear(2048,18)
model = model.to(device)
optimizer = optim.Adam(model.fc.parameters(), lr = 1e-3)
criterion = 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= criterion(output,label)
        loss.backward()
        optimizer.step()
    print(f"EPOCH{epoch+1} --- LOSS{loss:.4f}")   
    
class TestDataset(Dataset):
    def __init__(self, cv_path, transform):
        self.df = pd.read_csv(cv_path)
        self.transform = transform
    def __len__(self):
        return (len(self.df))
    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img= Image.open(row['ImagePath']).convert('RGB')
        if self.transform:
            img = self.transform(img)
        return img
test_ds = TestDataset("./test.csv", transform=transform)
test = DataLoader(test_ds, batch_size = 8, shuffle=False, num_workers=0)

pred = []
for img in test:
    with torch.no_grad():
        img = img.to(device)
        pred_per_file = model(img)
        _, predicted = torch.max(pred_per_file, 1)
        pred.append(predicted.cpu())
        preds = torch.cat(pred)
preds
preds = preds.numpy()
preds_rev = []
for i in preds:
    preds_rev.append(map_rev[i])
preds_rev
    
rows = []
for idx, row in test_df.iterrows():
    rows.append({"SampleID":int(row['ImagePath'].split("/")[1].strip(".png")),'Type':preds_rev[idx]})
pd.DataFrame(rows).to_csv('subs.csv',index=False)
    
