{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87323,"databundleVersionId":9934427,"sourceType":"competition"}],"dockerImageVersionId":30823,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.tensorboard import SummaryWriter\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport pandas as pd\nimport numpy as np\nimport cv2\nfrom tqdm import tqdm\nimport timm\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import train_test_split\n\n\nclass EfficientDataset(Dataset):\n    def __init__(self, df, img_dir, transform=None):\n        self.df = df\n        self.img_dir = img_dir\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = os.path.join(self.img_dir, row['image'])\n        \n        image = cv2.imread(img_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        \n        if 'label' in row:\n            label = torch.tensor(row['label'], dtype=torch.float32)\n            return image, label\n        return image,  \n\n\nclass LightweightModel(nn.Module):\n    def __init__(self, model_name='tf_efficientnet_b7', num_classes=1):\n        super(LightweightModel, self).__init__()\n        self.model_name = model_name\n        self.backbone = timm.create_model(model_name, pretrained=True)\n        \n        if hasattr(self.backbone, 'classifier'):\n            in_features = self.backbone.classifier.in_features\n            self.backbone.classifier = nn.Sequential(\n                nn.Dropout(0.4),  # Increased dropout\n                nn.Linear(in_features, num_classes)\n            )\n        else:\n            in_features = self.backbone.fc.in_features\n            self.backbone.fc = nn.Sequential(\n                nn.Dropout(0.4),  # Increased dropout\n                nn.Linear(in_features, num_classes)\n            )\n        \n    def forward(self, x):\n        return self.backbone(x)\n\n\ndef get_efficient_transforms(img_size):\n    train_transform = A.Compose([\n        A.Resize(img_size, img_size),\n        A.HorizontalFlip(p=0.5),\n        A.RandomBrightnessContrast(p=0.2),\n        A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=20, p=0.5),\n        A.CoarseDropout(max_holes=8, max_height=16, max_width=16, p=0.5),  # Replace Cutout\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n    \n    valid_transform = A.Compose([\n        A.Resize(img_size, img_size),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n    \n    return train_transform, valid_transform\n\n\ndef train_model(model, train_loader, val_loader, device, config):\n    criterion = nn.BCEWithLogitsLoss()\n    optimizer = optim.AdamW(model.parameters(), lr=config['learning_rate'], weight_decay=config['weight_decay'])\n    scheduler = optim.lr_scheduler.OneCycleLR(\n        optimizer,\n        max_lr=config['learning_rate'],\n        epochs=config['epochs'],\n        steps_per_epoch=len(train_loader),\n        div_factor=25,\n        final_div_factor=1e4\n    )\n    \n    best_auc = 0\n    early_stop_counter = 0\n    metrics = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': [], 'val_auc': []}\n    save_path = f\"best_model_{model.model_name}.pth\"\n    \n    writer = SummaryWriter(log_dir='logs')\n    \n    for epoch in range(config['epochs']):\n        model.train()\n        train_loss = 0\n        correct_train = 0\n        total_train = 0\n        \n        for images, labels in tqdm(train_loader, desc=f'Epoch {epoch + 1}/{config[\"epochs\"]}'):\n            images = images.to(device)\n            labels = labels.to(device).unsqueeze(1)\n            \n            optimizer.zero_grad()\n            outputs = model(images)\n            loss = criterion(outputs, labels)\n            \n            loss.backward()\n            optimizer.step()\n            scheduler.step()\n            \n            train_loss += loss.item()\n            preds = torch.sigmoid(outputs) > 0.5\n            correct_train += (preds == labels).sum().item()\n            total_train += labels.size(0)\n        \n        train_acc = correct_train / total_train\n\n        model.eval()\n        val_preds = []\n        val_labels = []\n        val_loss = 0\n        correct_val = 0\n        total_val = 0\n        \n        with torch.no_grad():\n            for images, labels in val_loader:\n                images = images.to(device)\n                labels = labels.to(device).unsqueeze(1)\n                \n                outputs = model(images)\n                loss = criterion(outputs, labels)\n                val_loss += loss.item()\n                \n                val_preds.extend(torch.sigmoid(outputs).cpu().numpy())\n                val_labels.extend(labels.cpu().numpy())\n                \n                preds = torch.sigmoid(outputs) > 0.5\n                correct_val += (preds == labels).sum().item()\n                total_val += labels.size(0)\n        \n        val_acc = correct_val / total_val\n        val_preds = np.array(val_preds)\n        val_labels = np.array(val_labels)\n        val_auc = roc_auc_score(val_labels, val_preds)\n        \n        metrics['train_loss'].append(train_loss / len(train_loader))\n        metrics['val_loss'].append(val_loss / len(val_loader))\n        metrics['train_acc'].append(train_acc)\n        metrics['val_acc'].append(val_acc)\n        metrics['val_auc'].append(val_auc)\n        writer.add_scalar('Loss/train', metrics['train_loss'][-1], epoch)\n        writer.add_scalar('Loss/val', metrics['val_loss'][-1], epoch)\n        writer.add_scalar('Accuracy/train', train_acc, epoch)\n        writer.add_scalar('Accuracy/val', val_acc, epoch)\n        writer.add_scalar('AUC/val', val_auc, epoch)\n        \n        print(f'\\nEpoch {epoch + 1}:')\n        print(f'Train Loss: {metrics[\"train_loss\"][-1]:.4f}, Train Accuracy: {train_acc:.4f}')\n        print(f'Val Loss: {metrics[\"val_loss\"][-1]:.4f}, Val Accuracy: {val_acc:.4f}, Val AUC: {val_auc:.4f}')\n        \n        if val_auc > best_auc:\n            best_auc = val_auc\n            early_stop_counter = 0\n            torch.save({\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict(),\n                'best_auc': best_auc,\n                'epoch': epoch\n            }, save_path)\n        else:\n            early_stop_counter += 1\n            if early_stop_counter >= config['patience']:\n                print(f\"Early stopping at epoch {epoch + 1}\")\n                break\n    \n    writer.close()\n    \n    metrics_df = pd.DataFrame(metrics)\n    metrics_df.to_csv('training_metrics.csv', index=False)\n    \n    return best_auc, save_path\n\n\ndef predict(model, test_loader, device, weights_path):\n    checkpoint = torch.load(weights_path, map_location=device)\n    model.load_state_dict(checkpoint['model_state_dict'])\n    \n    model.to(device)\n    model.eval()\n    predictions = []\n    \n    with torch.no_grad():\n        for images, _ in tqdm(test_loader, desc='Predicting'): \n            images = images.to(device)\n            outputs = model(images)\n            predictions.extend(torch.sigmoid(outputs).cpu().numpy())\n    \n    return np.array(predictions)\n\n\ndef main():\n    config = {\n        'img_size': 224,\n        'batch_size': 4,\n        'epochs': 50,\n        'learning_rate': 1e-3,\n        'patience': 5,\n        'weight_decay': 1e-4,\n        'train_val_split': 0.2,\n        'seed': 42\n    }\n\n    torch.manual_seed(config['seed'])\n    np.random.seed(config['seed'])\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    train_df = pd.read_csv(\"/kaggle/input/cidaut-ai-fake-scene-classification-2024/train.csv\")\n    train_df['label'] = train_df['label'].map({\"editada\": 0, \"real\": 1})\n    \n    train_data, val_data = train_test_split(train_df, test_size=config['train_val_split'], random_state=config['seed'])\n    train_transform, valid_transform = get_efficient_transforms(config['img_size'])\n    \n    train_dataset = EfficientDataset(\n        train_data,\n        \"/kaggle/input/cidaut-ai-fake-scene-classification-2024/Train\",\n        transform=train_transform\n    )\n    val_dataset = EfficientDataset(\n        val_data,\n        \"/kaggle/input/cidaut-ai-fake-scene-classification-2024/Train\",\n        transform=valid_transform\n    )\n    \n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=config['batch_size'],\n        shuffle=True,\n        num_workers=2\n    )\n    val_loader = DataLoader(\n        val_dataset,\n        batch_size=config['batch_size'],\n        shuffle=False,\n        num_workers=2\n    )\n    \n    model_name = 'tf_efficientnet_b7'\n    print(f\"\\nTraining {model_name}\")\n    model = LightweightModel(model_name=model_name).to(device)\n\n    score, weights_path = train_model(model, train_loader, val_loader, device, config)\n    print(f\"{model_name} - Best AUC: {score:.4f}\")\n    test_df = pd.read_csv(\"/kaggle/input/cidaut-ai-fake-scene-classification-2024/sample_submission.csv\")\n    test_dataset = EfficientDataset(\n        test_df,\n        \"/kaggle/input/cidaut-ai-fake-scene-classification-2024/Test\",\n        transform=valid_transform\n    )\n    test_loader = DataLoader(\n        test_dataset,\n        batch_size=config['batch_size'],\n        shuffle=False,\n        num_workers=2\n    )\n    \n    predictions = predict(model, test_loader, device, weights_path)\n    \n    test_df['label'] = predictions.flatten()\n    test_df.to_csv('submission.csv', index=False)\n    print(\"\\nPredictions saved to submission.csv\")\n\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-25T14:16:54.699532Z","iopub.execute_input":"2024-12-25T14:16:54.699849Z","iopub.status.idle":"2024-12-25T14:27:10.257089Z","shell.execute_reply.started":"2024-12-25T14:16:54.699825Z","shell.execute_reply":"2024-12-25T14:27:10.255849Z"}},"outputs":[],"execution_count":null}]}