

import pandas as pd
train = pd.read_csv("./train_data.csv")
train.head()
train['style'].unique()
import torch
import torch.nn as nn
from torchvision import transforms, models
from torch.utils.data import DataLoader,Dataset,random_split
from PIL import Image
import torch.optim as optim
map_style = {'cartoon':0, 'photo':1}
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'])
        label = map_style[row['style']]
        if self.transform:
            img = self.transform(img)
        return img,label
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
])
full_ds = TrainDataset("./train_data.csv",transform)
train_subset, val_subset = random_split(full_ds, [0.8,0.2])
train_loader = DataLoader(train_subset, shuffle=True, num_workers=0, batch_size=32)
val_loader = DataLoader(val_subset, shuffle=False, num_workers=0, batch_size=32)
device = "cuda"
from sklearn.metrics import f1_score
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(512,2)
for param in model.parameters():
    param.requires_grad = True
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="max")
criterion = nn.CrossEntropyLoss()
for epoch in range(20):
    model.train()
    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()
    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for img, label in val_loader:
            img=img.to(device)
            label = label.to(device)
            output = model(img)
            pred = output.argmax(dim=1)
            all_preds.append(pred.cpu())
            all_labels.append(label.cpu())
    all_labels = torch.cat(all_labels)
    all_preds = torch.cat(all_preds)
    all_preds = all_preds.numpy()
    all_labels = all_labels.numpy()
    f1 = f1_score(all_labels, all_preds, average='macro')
    scheduler.step(f1)
    print(f"Epoch {epoch+1}, loss {total_loss/len(train_loader)}, f1 {f1}")
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'])
        if self.transform:
            img = self.transform(img)
        return img
test_ds =TestDataset("./test_data.csv", transform)
test_loader= DataLoader(test_ds, shuffle=False, num_workers=0, batch_size=32)
preds = []
model.eval()
with torch.no_grad():
    for img in test_loader:
        img = img.to(device)
        output = model(img)
        pred = output.argmax(dim = 1)
        preds.append(pred.cpu())
preds_final = torch.cat(preds)
preds_final = preds_final.numpy()
preds_final
test = pd.read_csv('./test_data.csv')
test.head()
rev_map ={}
for key,values in map_style.items():
    rev_map[values] = key
pred_final_real = []
for i in preds_final:
    pred_final_real.append(rev_map[i])
pred_final_real
sub1 = pd.DataFrame({
    'subtaskID':1,
    'datapointID':test['ID'],
    'answer':pred_final_real
})
# STRATEGY: nn.Identity(), cosine similarity, daca sunt prea diferite atunci sunt 2 obiecte diferite, fac embeddings cu resnet pt ca e antrenat pe obiecte din viata reala
model2 = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
model2.fc = nn.Identity()
model2 = model2.to(device)
model2.eval()
import os
class Task2Dataset(Dataset):
    def __init__(self, folder_path, transform=None):
        self.folder_path = folder_path
        self.transform = transform
        
        self.image_paths = sorted(
            [
                os.path.join(folder_path, img)
                for img in os.listdir(folder_path)
                if img.endswith(".png")
            ],
            key=lambda x: int(os.path.splitext(os.path.basename(x))[0])
        )

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, index):
        img_path = self.image_paths[index]
        img = Image.open(img_path).convert("RGB")
        
        if self.transform:
            img = self.transform(img)
            
        return img
task2_ds = Task2Dataset("./images",transform)
task2_dl = DataLoader(task2_ds, shuffle=False, num_workers=0, batch_size=64)
embeddings = []
with torch.no_grad():
    for img in task2_dl:
        img = img.to(device)
        output = model2(img)
        embeddings.append(output.cpu())
embeddings = torch.cat(embeddings)
import numpy as np
from sklearn.cluster import KMeans

# embeddings: tensor deja calculat cu model2 (fc = Identity)
emb = embeddings.numpy()
emb = emb / (np.linalg.norm(emb, axis=1, keepdims=True) + 1e-12)

# predictiile de la task 1: lista cu "cartoon"/"photo"
style_preds = pred_final_real

cartoon_idx = [i for i, x in enumerate(style_preds) if x == "cartoon"]
photo_idx = [i for i, x in enumerate(style_preds) if x == "photo"]

# clustering separat
km_cartoon = KMeans(n_clusters=2, random_state=42, n_init=10)
cartoon_clusters = km_cartoon.fit_predict(emb[cartoon_idx])

km_photo = KMeans(n_clusters=2, random_state=42, n_init=10)
photo_clusters = km_photo.fit_predict(emb[photo_idx])

# vector final de clustere
cluster_labels = np.empty(len(emb), dtype=int)
for idx, c in zip(cartoon_idx, cartoon_clusters):
    cluster_labels[idx] = c
for idx, c in zip(photo_idx, photo_clusters):
    cluster_labels[idx] = c

# exemple etichetate manual
elephant_idx = [1, 2, 4, 11, 18, 19]
house_idx = [0, 3, 5, 6, 7, 9, 10]

final_labels = [""] * len(emb)

# mapping separat pe cartoon
known_cartoon_elephants = [i for i in elephant_idx if i in cartoon_idx]
known_cartoon_houses = [i for i in house_idx if i in cartoon_idx]

if len(known_cartoon_elephants) > 0 and len(known_cartoon_houses) > 0:
    elephant_cluster_cartoon = max(
        [cluster_labels[i] for i in known_cartoon_elephants],
        key=[cluster_labels[i] for i in known_cartoon_elephants].count
    )
    house_cluster_cartoon = 1 - elephant_cluster_cartoon

    for i in cartoon_idx:
        final_labels[i] = "elephant" if cluster_labels[i] == elephant_cluster_cartoon else "house"

# mapping separat pe photo
known_photo_elephants = [i for i in elephant_idx if i in photo_idx]
known_photo_houses = [i for i in house_idx if i in photo_idx]

if len(known_photo_elephants) > 0 and len(known_photo_houses) > 0:
    elephant_cluster_photo = max(
        [cluster_labels[i] for i in known_photo_elephants],
        key=[cluster_labels[i] for i in known_photo_elephants].count
    )
    house_cluster_photo = 1 - elephant_cluster_photo

    for i in photo_idx:
        final_labels[i] = "elephant" if cluster_labels[i] == elephant_cluster_photo else "house"

# fortam exemplele cunoscute
for i in elephant_idx:
    final_labels[i] = "elephant"
for i in house_idx:
    final_labels[i] = "house"
np.unique(final_labels)
train.head()
len(final_labels)
t = [i for i in range(1227)]
len(t)
test['ID']
final_labels_test = final_labels[980:1226]
final_labels_test
sub2 = pd.DataFrame({
    'subtaskID': 2,
    'datapointID': test['ID'],
    'answer': final_labels
})
sub_final = pd.concat([sub1,sub2])
sub_final.to_csv('subs.csv',index=False)
