

import pandas as pd
import numpy as np
df_train = pd.read_csv("train_data.csv")
df_test = pd.read_csv("test_data.csv")
df_train.head(5)
df_test = pd.read_csv("test_data.csv")
df_test.head(5)
from sklearn.preprocessing import LabelEncoder
lbl = LabelEncoder()
df_train["style"] = lbl.fit_transform(df_train["style"])
from torchvision.transforms import v2
train_transform = v2.Compose([
    v2.Resize((224, 224)),
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomVerticalFlip(p=0.2),
    v2.RandomRotation(degrees=15),
    v2.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),
    v2.RandomGrayscale(p=0.1),
    v2.ToTensor(),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform = v2.Compose([
    v2.Resize((224, 224)),
    v2.ToTensor(),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
class data_idk(Dataset):
    def __init__(self, df, test = False, transform = test_transform):
        super().__init__()
        self.X_path = df["image_path"]
        self.test = test
        self.transform = transform
        if self.test == False:
          self.y = torch.tensor(np.array(df["style"]))
    def __len__(self):
        return len(self.X_path)
    def __getitem__(self, index):
        img = self.transform(Image.open(self.X_path[index]).convert("RGB"))
        if self.test == False:
            return img, self.y[index].float()
        return img
from torch.utils.data import random_split
full_dt = data_idk(df_train, transform= train_transform)
train_len = int(0.8 * len(full_dt))
valid_len = len(full_dt) - train_len
train_dt, valid_dt = random_split(full_dt, [train_len, valid_len])
test_dt =data_idk(df_test, test= True)
train_loader = DataLoader(train_dt, batch_size= 32, shuffle= True)
valid_loader = DataLoader(valid_dt, batch_size= 32, shuffle= True)
test_loader = DataLoader(test_dt, batch_size= 32)

import torchvision
import torch.nn as nn
class model_idk(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.backbone = torchvision.models.resnet50(pretrained = True)
        self.in_feat = self.backbone.fc.in_features
        self.backbone.fc = nn.Identity()
        self.cls = nn.Sequential(
            nn.Linear(self.in_feat, 512),
            nn.ReLU(),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
        )
    def forward(self, x):
        x = self.backbone(x)
        x = self.cls(x)
        return x.squeeze(1)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model = model_idk()
model = model.to(DEVICE)
criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters())
NUM_EPOCHS = 10
from sklearn.metrics import f1_score
from tqdm.notebook import tqdm

history = {'train_loss': [], 'train_f1': [], 'val_loss': [], 'val_f1': []}
best_val_f1 = 0.0

for epoch in range(1, NUM_EPOCHS + 1):

    # ── Train ──
    model.train()
    train_loss, train_preds, train_targets = 0.0, [], []

    pbar = tqdm(train_loader, desc=f'Epoch {epoch:02d}/{NUM_EPOCHS} [train]', leave=False)
    for imgs, labels in pbar:
        imgs   = imgs.to(DEVICE)
        labels = labels.to(DEVICE).float()

        optimizer.zero_grad()
        logits = model(imgs)         # [B]
        loss   = criterion(logits, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item() * imgs.size(0)
        preds = (logits.sigmoid() > 0.5).long().cpu().tolist()
        train_preds.extend(preds)
        train_targets.extend(labels.long().cpu().tolist())

        pbar.set_postfix(loss=f'{loss.item():.4f}')

    train_loss /= len(train_loader.dataset)
    train_f1    = f1_score(train_targets, train_preds, average='binary')

    # ── Validate ──
    model.eval()
    val_loss, val_preds, val_targets = 0.0, [], []

    with torch.no_grad():
        pbar = tqdm(valid_loader, desc=f'Epoch {epoch:02d}/{NUM_EPOCHS} [val]  ', leave=False)
        for imgs, labels in pbar:
            imgs   = imgs.to(DEVICE)
            labels = labels.to(DEVICE).float()
            logits = model(imgs)
            loss   = criterion(logits, labels)

            val_loss += loss.item() * imgs.size(0)
            preds = (logits.sigmoid() > 0.5).long().cpu().tolist()
            val_preds.extend(preds)
            val_targets.extend(labels.long().cpu().tolist())

            pbar.set_postfix(loss=f'{loss.item():.4f}')

    val_loss /= len(valid_loader.dataset)
    val_f1    = f1_score(val_targets, val_preds, average='binary')

    # ── Log ──
    history['train_loss'].append(train_loss)
    history['train_f1'].append(train_f1)
    history['val_loss'].append(val_loss)
    history['val_f1'].append(val_f1)

    flag = ''
    if val_f1 > best_val_f1:
        best_val_f1 = val_f1
        torch.save(model.state_dict(), 'best_model.pth')
        flag = '  ← best'

    print(
        f'Epoch {epoch:02d}/{NUM_EPOCHS}  '
        f'train_loss={train_loss:.4f}  train_f1={train_f1:.4f}  '
        f'val_loss={val_loss:.4f}  val_f1={val_f1:.4f}'
        + flag
    )

print(f'\nBest val F1: {best_val_f1:.4f}')

model.load_state_dict(torch.load("best_model.pth"))
print("ALL WORKS")
test_preds = []
for batch in tqdm(test_loader, total= len(test_loader)):
    X = batch.to(DEVICE)
    y_pred =(torch.sigmoid(model(X)) > 0.5).int().cpu().numpy().tolist()
    test_preds.extend(y_pred)
df_pred_1 = pd.DataFrame({
    "subtaskID" : 1, 
    "datapointID" : df_test["ID"],
    "answer" : lbl.inverse_transform(test_preds)
})
df_test["style"] = lbl.inverse_transform(test_preds)
df_test.head(5)
model = torchvision.models.resnet50(pretrained = True)
model.fc = nn.Identity()
print("OK")
df_test_cartoon = df_test[df_test["style"] == "cartoon"]
with torch.no_grad():
    embeddings = []
    for i in tqdm(range(len(df_test_cartoon))):
        img = Image.open(df_test_cartoon["image_path"].iloc[i])
        img = transform(img)
        img = model(img.unsqueeze(0))
        embeddings.append(img.squeeze(0))
    embeddings = torch.stack(embeddings)
    embeddings.shape
from sklearn.cluster import KMeans
kmn = KMeans(n_clusters= 2)
y_pred_cartoons = kmn.fit_predict(embeddings.detach().numpy())
df_test_cartoon["answer"] = y_pred_cartoons
df_test_photo = df_test[df_test["style"] == "photo"]
with torch.no_grad():
    embeddings = []
    for i in tqdm(range(len(df_test_photo))):
        img = Image.open(df_test_photo["image_path"].iloc[i])
        img = transform(img)
        img = model(img.unsqueeze(0))
        embeddings.append(img.squeeze(0))
    embeddings = torch.stack(embeddings)
    embeddings.shape
from sklearn.cluster import KMeans
kmn = KMeans(n_clusters= 2)
y_pred_photo = kmn.fit_predict(embeddings.detach().numpy())
df_test_photo["answer"] = y_pred_photo
df_test_photo["answer"] = df_test_photo["answer"].map({
    0 : 1,
    1: 0
})
df_test = pd.concat([df_test_cartoon, df_test_photo])
df_test.head(5)
df_pred_2 = pd.DataFrame({
    "subtaskID" : 2, 
    "datapointID" : df_test["ID"],
    "answer" : df_test["answer"]
})
df_pred = pd.concat([df_pred_1, df_pred_2])
df_pred.to_csv("subi.csv", index= False)
