{"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":"nvidiaTeslaT4","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\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\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='efficientnet_lite0', num_classes=1):\n        super(LightweightModel, self).__init__()\n        self.model_name = model_name  # Save model name for reference\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.3),\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.3),\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.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=0.01)\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    )\n    \n    best_auc = 0\n\n    save_path = f\"best_model_{model.model_name}.pth\"\n    \n    for epoch in range(config['epochs']):\n\n        model.train()\n        train_loss = 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        \n        model.eval()\n        val_preds = []\n        val_labels = []\n        val_loss = 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        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        print(f'Epoch {epoch + 1}:')\n        print(f'Train Loss: {train_loss/len(train_loader):.4f}')\n        print(f'Val Loss: {val_loss/len(val_loader):.4f}')\n        print(f'Val AUC: {val_auc:.4f}')\n\n        if val_auc > best_auc:\n            best_auc = val_auc\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    \n    return best_auc, save_path\n\n\ndef predict(model, test_loader, device, weights_path):\n\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': 32,\n        'epochs': 15,\n        'learning_rate': 1e-3,\n        'patience': 3,\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\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    val_mask = np.random.rand(len(train_df)) < config['train_val_split']\n    train_data = train_df[~val_mask]\n    val_data = train_df[val_mask]\n    \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 = 'efficientnet_lite0'\n    print(f\"\\nTraining {model_name}\")\n    model = LightweightModel(model_name=model_name).to(device)\n    score, weights_path = train_model(model, train_loader, val_loader, device, config)\n    print(f\"{model_name} - Best AUC: {score:.4f}\")\n\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()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-12-21T09:22:23.981976Z","iopub.execute_input":"2024-12-21T09:22:23.982270Z","iopub.status.idle":"2024-12-21T09:23:43.088262Z","shell.execute_reply.started":"2024-12-21T09:22:23.982249Z","shell.execute_reply":"2024-12-21T09:23:43.087304Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}