{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":129617,"databundleVersionId":15872846}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Welcome to the ACM AI Mini ML Challenge!\nIn this challenge, your task is to build a machine learning model capable of classifying Star Wars characters as either \"Good\" or \"Bad\" (Light side vs. Dark side).\n\nBecause we are working with complex image data, training a model from scratch would take hours (or days!) and require millions of images. Instead, we will use a technique called **Transfer Learning**.\n\nThis starter notebook will help you:\n1. Load and visualize the dataset.\n2. Understand the core concepts of Transfer Learning and Hyperparameters.\n3. Build the scaffolding for your PyTorch model.\n4. Format your final predictions for submission.","metadata":{}},{"cell_type":"code","source":"# ==========================================\n# 1. SETUP & IMPORTS\n# ==========================================\n# Install required libraries if you are running locally:\n# %pip install numpy pandas pillow torch torchvision matplotlib\n\nimport os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchvision import transforms\nimport torchvision.models as models\nimport torchvision.transforms.functional as F2\n\n# Challenge Dataset Paths\ntrain_path = \"/kaggle/input/competitions/acm-ai-mini-ml-challenge-2026/train/train\"\ntest_path = \"/kaggle/input/competitions/acm-ai-mini-ml-challenge-2026/test/test\"\ncsv_drive_path = \"/kaggle/input/competitions/acm-ai-mini-ml-challenge-2026\"\n\n# Set device to GPU if available, otherwise CPU\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T22:21:50.056438Z","iopub.execute_input":"2026-02-28T22:21:50.057253Z","iopub.status.idle":"2026-02-28T22:21:50.063330Z","shell.execute_reply.started":"2026-02-28T22:21:50.057223Z","shell.execute_reply":"2026-02-28T22:21:50.062754Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Part 1: Loading the Dataset\nWe use a custom PyTorch `Dataset` class to load our images and map them to their labels. The data is provided in CSV format, linking image file paths to character names and their target alignment (Good=1, Bad=0). \n\n**Important:** Our paths in `train.csv` already include the `train/` folder prefix, so we set our `root_dir` simply to `\".\"` (the current directory).","metadata":{}},{"cell_type":"code","source":"# ==========================================\n# 2. DATASET DEFINITION\n# ==========================================\n\nclass ImageCSVDataset(Dataset):\n    def __init__(self, csv_file, root_dir, transform=None, label_map=None, is_test=False):\n        self.df = pd.read_csv(csv_file)\n        self.root_dir = root_dir\n        self.label_map = label_map\n        self.transform = transform\n        self.is_test = is_test\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        \n        # We use os.path.basename to clean up the path just in case \n        # the CSV contains \"train/img_0.jpg\" so we don't end up with \"train/train/img_0.jpg\"\n        clean_filename = os.path.basename(row['path']) \n        img_path = os.path.join(self.root_dir, clean_filename)\n\n        # Convert to RGB to ensure 3 channels for pre-trained models\n        image = Image.open(img_path).convert('RGB')\n\n        if self.transform:\n            image = self.transform(image)\n\n        if self.is_test:\n            image_id = int(row['ID'])\n            return image, image_id  # useful for mapping back to submission\n        else:\n            label = self.label_map[row['target']]\n            character = row['character']\n            return image, label, character\n\n# Define standard ImageNet transformations\ntransform = transforms.Compose([\n    transforms.Resize((224, 224)),\n    transforms.RandomHorizontalFlip(), # 随机水平翻转\n    transforms.RandomRotation(15),      # 随机旋转\n    transforms.ColorJitter(brightness=0.1, contrast=0.1), # 亮度对比度微调\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    transforms.RandomHorizontalFlip(),\n    transforms.ColorJitter()\n])\n\n\n\nlabel_map = {\"bad\": 0, \"good\": 1}\n\n# Instantiate the training dataset\ntrain_dataset = ImageCSVDataset(\n    csv_file=os.path.join(csv_drive_path, \"train.csv\"), \n    root_dir=train_path,    \n    transform=transform,\n    label_map=label_map,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T22:21:50.064666Z","iopub.execute_input":"2026-02-28T22:21:50.064898Z","iopub.status.idle":"2026-02-28T22:21:50.084065Z","shell.execute_reply.started":"2026-02-28T22:21:50.064867Z","shell.execute_reply":"2026-02-28T22:21:50.083442Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Part 2: Data Playground\nBefore training any model, it's critical to perform Exploratory Data Analysis (EDA). This helps us understand what the model will \"see\". Run the cell below to visualize the class balance and inspect random samples from the training set. Feel free to re-run it multiple times!","metadata":{}},{"cell_type":"code","source":"# ==========================================\n# 3. VISUALIZATION PLAYGROUND\n# ==========================================\n\n# 1. Analyze Class Distribution\nprint(\"--- Class Distribution ---\")\nclass_counts = train_dataset.df['target'].value_counts()\nprint(class_counts)\n\n# Plot a simple bar chart\nplt.figure(figsize=(6, 4))\nclass_counts.plot(kind='bar', color=['skyblue', 'salmon'])\nplt.title('Good vs. Bad Characters')\nplt.ylabel('Count')\nplt.xticks(rotation=0) # Keeps the text upright\nplt.show()\n\n# ------------------------------------------\n\n# 2. Display a Random Grid of Images\nprint(\"\\n--- Sample Images ---\")\nplt.figure(figsize=(10, 10))\n\n# Grab 9 random rows from our dataframe directly\nsample_df = train_dataset.df.sample(9)\n\n# Loop through those 9 rows and plot them\nfor i, (_, row) in enumerate(sample_df.iterrows()):\n    # Reconstruct the exact path using our Kaggle root path\n    # (Assumes your dataframe paths look like 'train/img_0.jpg')\n    img_path = os.path.join(\"/kaggle/input/competitions/acm-ai-mini-ml-challenge-2026/train\", row['path'])\n    \n    # Open the raw image (no math required!)\n    img = Image.open(img_path)\n    \n    # Plot it in a 3x3 grid\n    plt.subplot(3, 3, i + 1)\n    plt.imshow(img)\n    \n    # Add a nice title\n    label_name = \"Good\" if row['target'] == \"good\" else \"Bad\"\n    plt.title(f\"{row['character']}\\n({label_name})\", fontsize=10)\n    plt.axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T22:21:50.085114Z","iopub.execute_input":"2026-02-28T22:21:50.085357Z","iopub.status.idle":"2026-02-28T22:21:52.278499Z","shell.execute_reply.started":"2026-02-28T22:21:50.085331Z","shell.execute_reply":"2026-02-28T22:21:52.277466Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Part 3: Core Concepts\n\n### What is Transfer Learning?\n\n\nTransfer Learning is the practice of taking a model that has already been trained on a massive dataset (like ImageNet, which contains millions of images) and repurposing it for a specific task—in our case, classifying Star Wars characters. \n\nThe pre-trained model has already learned how to detect basic visual features like edges, shapes, and textures. Instead of training a model from scratch, we keep the model's foundational knowledge and simply \"swap out\" its final output layer (the \"head\") to output our specific classes (Good vs. Bad).\n\n### Tuning the Engine: Hyperparameters\nHyperparameters are the settings you configure *before* the model starts training. Think of them like the dials on a dashboard. Modifying these will dictate how well your model learns:\n\n* **Epochs:** How many times the model will look at the entire dataset. Too few, and it won't learn enough (underfitting). Too many, and it will memorize the training data but fail on new data (overfitting).\n* **Batch Size:** How many images the model processes at once before updating its knowledge. Common values are 16, 32, or 64.\n* **Learning Rate (LR):** The \"step size\" the model takes when adjusting its weights. \n    * *High LR:* The model learns fast but might overshoot the optimal solution.\n    * *Low LR:* The model learns carefully but very slowly.\n* **Optimizer:** The algorithm used to update the model (e.g., `Adam` or `SGD`).","metadata":{}},{"cell_type":"code","source":"# ==========================================\n# 4. MODEL DEFINITION & HYPERPARAMETERS\n# ==========================================\n\nclass StarWarsClassifier(nn.Module):\n    def __init__(self, num_classes=2):\n        super(StarWarsClassifier, self).__init__()\n        \n        # Load a pre-trained ResNet18 model\n        self.model = models.resnet18(pretrained=True)\n\n        # Optional: Freeze the early layers so they don't update during training.\n        # This speeds up training and prevents the model from forgetting basic features.\n        # for param in self.model.parameters():\n        #     param.requires_grad = False\n\n        # Replace the final fully connected layer of the ResNet18 model\n        self.model.fc = nn.Linear(self.model.fc.in_features, 2)\n        \n    def forward(self, x):\n        return self.model(x)\n\n# Initialize Model, Loss Function, and Optimizer\nmodel = StarWarsClassifier(num_classes=2).to(device)\ncriterion = nn.CrossEntropyLoss()\n\n# --- PLAY WITH THESE HYPERPARAMETERS ---\nLEARNING_RATE = 0.0005\nBATCH_SIZE = 32\nEPOCHS = 5\n# ---------------------------------------\n\noptimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\nscheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T22:21:52.280145Z","iopub.execute_input":"2026-02-28T22:21:52.280412Z","iopub.status.idle":"2026-02-28T22:21:52.473521Z","shell.execute_reply.started":"2026-02-28T22:21:52.280383Z","shell.execute_reply":"2026-02-28T22:21:52.472779Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Part 4: Training the Model\nWe split our training data into a `train` set and a `validation` set. The validation set allows us to test the model on images it hasn't seen during training, giving us an idea of how well it generalizes.","metadata":{}},{"cell_type":"code","source":"# ==========================================\n# 5. THE TRAINING LOOP\n# ==========================================\n\n# Split data 80/20 for training and validation\ntrain_size = int(0.8 * len(train_dataset))\nval_size = len(train_dataset) - train_size\ntrain_subset, val_subset = random_split(train_dataset, [train_size, val_size])\n\ntrain_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True)\nval_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False)\n\ndef train_one_epoch(model, dataloader):\n    model.train()\n    total_loss = 0\n\n    for batch in dataloader:\n        imgs, labels, _ = batch\n        imgs, labels = imgs.to(device), labels.to(device)\n\n        optimizer.zero_grad()               # 1. 清空旧梯度\n        outputs = model(imgs)               # 2. 前向传播（得到预测结果）\n        loss = criterion(outputs, labels)    # 3. 计算误差\n        loss.backward()                     # 4. 反向传播（计算梯度）\n        optimizer.step()                    # 5. 更新模型参数（学习）\n        \n        total_loss += loss.item()\n\n    return total_loss / len(dataloader)\n\ndef evaluate(model, dataloader):\n    model.eval()\n    correct, total = 0, 0\n\n    with torch.no_grad(): # 评估时不需要计算梯度，节省内存\n        for batch in dataloader:\n            imgs, labels, _ = batch\n            imgs, labels = imgs.to(device), labels.to(device)\n            \n            outputs = model(imgs)\n            # torch.max 返回 (最大值, 索引)，索引即为预测的类别 (0 或 1)\n            _, predicted = torch.max(outputs.data, 1)\n            \n            total += labels.size(0)\n            correct += (predicted == labels).sum().item()\n            \n    return correct / total if total > 0 else 0.0\n\n# Start Training!\nprint(\"Starting Training...\")\nfor epoch in range(EPOCHS):\n    train_loss = train_one_epoch(model, train_loader)\n    val_acc = evaluate(model, val_loader)\n    scheduler.step()\n    print(f\"Epoch {epoch+1}/{EPOCHS}: Train Loss = {train_loss:.4f} | Validation Acc = {val_acc:.2%}\")\n    if val_acc > best_acc:\n        best_acc = val_acc\n        torch.save(model.state_dict(), 'best_model.pth')\n        print(f\"Better Model Saved! Acc: {val_acc:.2%}\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T22:21:52.474308Z","iopub.execute_input":"2026-02-28T22:21:52.474580Z","iopub.status.idle":"2026-02-28T22:24:04.695412Z","shell.execute_reply.started":"2026-02-28T22:21:52.474557Z","shell.execute_reply":"2026-02-28T22:24:04.694744Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics import confusion_matrix, classification_report\nimport seaborn as sns\n\n# 1. 获取验证集的所有真实标签和预测标签\nall_preds = []\nall_labels = []\n\nmodel.eval()\nwith torch.no_grad():\n    for imgs, labels, _ in val_loader:\n        imgs = imgs.to(device)\n        outputs = model(imgs)\n        _, predicted = torch.max(outputs, 1)\n        \n        all_preds.extend(predicted.cpu().numpy())\n        all_labels.extend(labels.cpu().numpy())\n\n# 2. 计算混淆矩阵\ncm = confusion_matrix(all_labels, all_preds)\n# label_map 是 {\"bad\": 0, \"good\": 1}，所以顺序是 Bad, Good\nclass_names = ['Bad (Dark Side)', 'Good (Light Side)']\n\n# 3. 绘图\nplt.figure(figsize=(8, 6))\nsns.heatmap(cm, annot=True, fmt='d', cmap='Blues', \n            xticklabels=class_names, yticklabels=class_names)\nplt.xlabel('Predicted (模型预测)')\nplt.ylabel('Actual (实际真值)')\nplt.title('Star Wars Character Classification - Confusion Matrix')\nplt.show()\n\n# 4. 打印详细报告（包含 Precision, Recall, F1-score）\nprint(\"\\n--- Detailed Classification Report ---\")\nprint(classification_report(all_labels, all_preds, target_names=class_names))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T22:41:39.755609Z","iopub.execute_input":"2026-02-28T22:41:39.755914Z","iopub.status.idle":"2026-02-28T22:41:46.425777Z","shell.execute_reply.started":"2026-02-28T22:41:39.755891Z","shell.execute_reply":"2026-02-28T22:41:46.424945Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ==========================================\n# 5.5 验证并锁定最佳模型权重 (在 Part 6 之前运行)\n# ==========================================\n\n# 1. 确保模型结构已经定义 (StarWarsClassifier)\n# 2. 这里的 model 变量就是你训练了 5 个 Epoch 的那个对象\n\n# 保存刚才 98.39% 的成果\nmodel_path = 'starwars_model_best.pth'\ntorch.save(model.state_dict(), model_path)\nprint(f\"权重已保存至 {model_path}\")\n\n# 3. 重新加载权重（这一步是为了确保文件没坏，且强行同步状态）\nmodel.load_state_dict(torch.load(model_path))\nmodel.to(device)\n\n# 4. 切换到评估模式 (这是衔接你的 Part 6 最关键的一步)\nmodel.eval()\n\nprint(\"验证加载成功！当前模型已锁定在 98.39% 的巅峰状态。\")\n\n# 5. (可选) 再次确认验证集准确率，确保万无一失\nval_acc = evaluate(model, val_loader)\nprint(f\"重新加载后的验证集准确率: {val_acc:.2%}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T22:40:26.438537Z","iopub.execute_input":"2026-02-28T22:40:26.438873Z","iopub.status.idle":"2026-02-28T22:40:31.958429Z","shell.execute_reply.started":"2026-02-28T22:40:26.438847Z","shell.execute_reply":"2026-02-28T22:40:31.957756Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Part 5: Generating your Submission\nOnce you are happy with your validation accuracy, run your model on the test set and save the results to a CSV file to submit to the leaderboard!","metadata":{}},{"cell_type":"code","source":"# ==========================================\n# 6. MAKE PREDICTIONS ON TEST DATA\n# ==========================================\n\ntest_dataset = ImageCSVDataset(\n    csv_file=os.path.join(csv_drive_path, \"test.csv\"), \n    root_dir=test_path,    \n    transform=transform,\n    is_test=True\n)\ntest_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n# Inverse map to turn integers (0, 1) back into strings (\"bad\", \"good\")\nidx2label = {v: k for k, v in label_map.items()}\n\nmodel.eval()\npreds = []\nids = []\n\nwith torch.no_grad():\n    for image, image_id in test_loader:\n        outputs = model(image.to(device))\n        _, predicted = torch.max(outputs, 1)\n        preds.extend(predicted.cpu().numpy())\n        ids.extend([id.item() if isinstance(id, torch.Tensor) else id for id in image_id])\n\npred_labels = [idx2label[p.item()] for p in preds]\n\nsubmission = pd.DataFrame({\n    \"ID\": ids,\n    \"target\": pred_labels\n})\n\nsubmission.to_csv(\"submission.csv\", index=False)\nprint(f\"submission.csv successfully created! Ready for upload.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-28T22:24:04.696401Z","iopub.execute_input":"2026-02-28T22:24:04.696694Z","iopub.status.idle":"2026-02-28T22:24:11.531174Z","shell.execute_reply.started":"2026-02-28T22:24:04.696662Z","shell.execute_reply":"2026-02-28T22:24:11.530558Z"}},"outputs":[],"execution_count":null}]}