{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"gpuType":"T4"},"kernelspec":{"name":"python3","display_name":"Python 3"},"accelerator":"GPU"},"cells":[{"cell_type":"markdown","source":"# Galaxy Zoo 5-Class Classification - Advanced Version\n## EfficientNet-B0 + Class Weighting + Advanced Augmentation\n\n**Improvements over baseline:**\n- EfficientNet-B0 architecture (better accuracy/params ratio)\n- SGD with momentum + cosine annealing\n- Class weighting for imbalanced classes\n- Advanced data augmentation\n- 50 epochs with early stopping\n- Test-time augmentation\n\n**Expected accuracy:** 85-90%","metadata":{"id":"header"}},{"cell_type":"markdown","source":"## Cell 1: Setup & Data Download","metadata":{"id":"cell1_header"}},{"cell_type":"code","source":"import os\nimport zipfile\n\n# Kaggle setup\nINPUT_DIR = '/kaggle/input/galaxy-zoo-the-galaxy-challenge'\nOUTPUT_DIR = '/kaggle/working'\n\ndef extract_data():\n    if not os.path.exists(INPUT_DIR):\n        print(f\"ERROR: Add dataset via Kaggle UI\")\n        return\n    \n    files = ['images_training_rev1.zip', 'training_solutions_rev1.zip']\n    for f in files:\n        path = os.path.join(INPUT_DIR, f)\n        if os.path.exists(path):\n            print(f\"Extracting {f}...\")\n            with zipfile.ZipFile(path, 'r') as z:\n                z.extractall(OUTPUT_DIR)\n    print(\"Done!\")\n\nextract_data()","metadata":{"id":"cell1"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 2: Data Loader (5 Classes)","metadata":{"id":"cell2_header"}},{"cell_type":"code","source":"%%writefile data_loader.py\nimport pandas as pd\nimport os\nfrom PIL import Image\nimport torch\nfrom torch.utils.data import Dataset\nimport numpy as np\n\nclass GalaxyZooDataset5Class(Dataset):\n    CLASS_NAMES = ['Round', 'In-between', 'Cigar', 'Edge-on', 'Spiral']\n    \n    def __init__(self, csv_file, root_dir, transform=None):\n        self.labels_df = pd.read_csv(csv_file)\n        self.root_dir = root_dir\n        self.transform = transform\n        self.labels = self._compute_labels()\n    \n    def _compute_labels(self):\n        labels = []\n        for idx in range(len(self.labels_df)):\n            row = self.labels_df.iloc[idx]\n            p_smooth = row['Class1.1']\n            p_features = row['Class1.2']\n            p_artifact = row['Class1.3']\n            \n            if p_artifact > 0.5:\n                labels.append(-1)\n                continue\n            \n            if p_smooth > p_features:\n                roundness = np.argmax([row['Class7.1'], row['Class7.2'], row['Class7.3']])\n                labels.append(roundness)\n            else:\n                if row['Class2.1'] > row['Class2.2']:\n                    labels.append(3)\n                else:\n                    labels.append(4)\n        return labels\n    \n    def __len__(self):\n        return len(self.labels_df)\n    \n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n        \n        label = self.labels[idx]\n        if label == -1:\n            return self.__getitem__((idx + 1) % len(self))\n        \n        img_id = str(self.labels_df.iloc[idx, 0])\n        img_name = os.path.join(self.root_dir, img_id + \".jpg\")\n        \n        try:\n            image = Image.open(img_name).convert('RGB')\n        except:\n            return self.__getitem__((idx + 1) % len(self))\n        \n        if self.transform:\n            image = self.transform(image)\n        \n        return image, int(label)","metadata":{"id":"cell2"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 3: Compute Class Weights","metadata":{"id":"cell3_header"}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom collections import Counter\n\n# Load labels to compute class weights\ndf = pd.read_csv('training_solutions_rev1.csv')\n\n# Compute labels\nlabels = []\nfor idx in range(len(df)):\n    row = df.iloc[idx]\n    p_smooth = row['Class1.1']\n    p_features = row['Class1.2']\n    p_artifact = row['Class1.3']\n    \n    if p_artifact > 0.5:\n        continue\n    \n    if p_smooth > p_features:\n        roundness = np.argmax([row['Class7.1'], row['Class7.2'], row['Class7.3']])\n        labels.append(roundness)\n    else:\n        if row['Class2.1'] > row['Class2.2']:\n            labels.append(3)\n        else:\n            labels.append(4)\n\n# Count per class\nclass_counts = Counter(labels)\nprint(\"Class distribution:\")\nclass_names = ['Round', 'In-between', 'Cigar', 'Edge-on', 'Spiral']\nfor i, name in enumerate(class_names):\n    print(f\"  {name}: {class_counts[i]} ({100*class_counts[i]/len(labels):.1f}%)\")\n\n# Compute weights (inverse frequency)\ntotal = len(labels)\nnum_classes = 5\nclass_weights = []\nfor i in range(num_classes):\n    weight = total / (num_classes * class_counts[i])\n    class_weights.append(weight)\n\nprint(f\"\\nClass weights: {class_weights}\")","metadata":{"id":"cell3"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 4: Setup Model & Data Loaders","metadata":{"id":"cell4_header"}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, Subset\nfrom torchvision import transforms, models\nfrom data_loader import GalaxyZooDataset5Class\nimport time\n\n# Hyperparameters\nBATCH_SIZE = 32\nLEARNING_RATE = 0.01  # Higher for SGD\nEPOCHS = 50\nNUM_CLASSES = 5\nPATIENCE = 10  # Early stopping patience\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nprint(f\"Training on: {DEVICE}\")\n\n# ADVANCED Data Augmentation\ntrain_transform = transforms.Compose([\n    transforms.Resize((240, 240)),  # Slightly larger for cropping\n    transforms.RandomCrop(224),\n    transforms.RandomRotation(360),  # Full rotation - galaxies are rotationally invariant\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomVerticalFlip(),\n    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    transforms.RandomErasing(p=0.1)  # Cutout-like augmentation\n])\n\nval_transform = transforms.Compose([\n    transforms.Resize((224, 224)),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n])\n\n# Load data\ndf = pd.read_csv('training_solutions_rev1.csv')\ndataset_length = len(df)\nindices = list(range(dataset_length))\n\n# Shuffle for better split\nimport random\nrandom.seed(42)\nrandom.shuffle(indices)\n\nsplit = int(0.8 * dataset_length)\ntrain_indices = indices[:split]\nval_indices = indices[split:]\n\ntrain_dataset_full = GalaxyZooDataset5Class(\n    csv_file='training_solutions_rev1.csv',\n    root_dir='images_training_rev1',\n    transform=train_transform\n)\n\nval_dataset_full = GalaxyZooDataset5Class(\n    csv_file='training_solutions_rev1.csv',\n    root_dir='images_training_rev1',\n    transform=val_transform\n)\n\ntrain_dataset = Subset(train_dataset_full, train_indices)\nval_dataset = Subset(val_dataset_full, val_indices)\n\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True)\n\nprint(f\"Training samples: {len(train_dataset)}\")\nprint(f\"Validation samples: {len(val_dataset)}\")\n\n# EfficientNet-B0 Model\nmodel = models.efficientnet_b0(weights='DEFAULT')\nmodel.classifier[1] = nn.Linear(model.classifier[1].in_features, NUM_CLASSES)\nmodel = model.to(DEVICE)\n\n# Print model info\ntotal_params = sum(p.numel() for p in model.parameters())\nprint(f\"\\nEfficientNet-B0 loaded!\")\nprint(f\"Total parameters: {total_params/1e6:.2f}M\")\n\n# Loss with class weights\nweights_tensor = torch.tensor(class_weights, dtype=torch.float32).to(DEVICE)\ncriterion = nn.CrossEntropyLoss(weight=weights_tensor)\n\n# SGD with momentum (better for fine-tuning)\noptimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=0.9, weight_decay=1e-4)\n\n# Cosine Annealing scheduler\nscheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-6)\n\nprint(\"\\nTraining setup complete!\")\nprint(f\"  Optimizer: SGD (lr={LEARNING_RATE}, momentum=0.9)\")\nprint(f\"  Scheduler: CosineAnnealing\")\nprint(f\"  Loss: CrossEntropy with class weights\")","metadata":{"id":"cell4"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 5: Training Functions with Early Stopping","metadata":{"id":"cell5_header"}},{"cell_type":"code","source":"def train_one_epoch(epoch):\n    model.train()\n    running_loss = 0.0\n    correct = 0\n    total = 0\n    \n    for i, (inputs, labels) in enumerate(train_loader):\n        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)\n        \n        optimizer.zero_grad()\n        outputs = model(inputs)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n        \n        running_loss += loss.item()\n        _, predicted = torch.max(outputs.data, 1)\n        total += labels.size(0)\n        correct += (predicted == labels).sum().item()\n        \n        if (i + 1) % 200 == 0:\n            print(f\"  Batch {i+1}/{len(train_loader)} | Loss: {loss.item():.4f} | Acc: {100*correct/total:.2f}%\")\n    \n    return running_loss / len(train_loader), 100 * correct / total\n\n\ndef validate():\n    model.eval()\n    running_loss = 0.0\n    correct = 0\n    total = 0\n    \n    with torch.no_grad():\n        for inputs, labels in val_loader:\n            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)\n            outputs = model(inputs)\n            loss = criterion(outputs, labels)\n            \n            running_loss += loss.item()\n            _, predicted = torch.max(outputs.data, 1)\n            total += labels.size(0)\n            correct += (predicted == labels).sum().item()\n    \n    return running_loss / len(val_loader), 100 * correct / total\n\nprint(\"Training functions defined with early stopping support!\")","metadata":{"id":"cell5"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 6: Main Training Loop with Early Stopping","metadata":{"id":"cell6_header"}},{"cell_type":"code","source":"# Training loop\nstart_time = time.time()\ntrain_losses, val_losses = [], []\ntrain_accs, val_accs = [], []\n\nbest_val_loss = float('inf')\nbest_val_acc = 0.0\nno_improve_count = 0\n\nprint(f\"🚀 Starting training for up to {EPOCHS} epochs...\")\nprint(f\"   Early stopping patience: {PATIENCE} epochs\")\nprint(\"=\"*60)\n\nfor epoch in range(EPOCHS):\n    current_lr = optimizer.param_groups[0]['lr']\n    print(f\"\\n--- Epoch {epoch+1}/{EPOCHS} (LR: {current_lr:.6f}) ---\")\n    \n    # Train\n    t_loss, t_acc = train_one_epoch(epoch)\n    train_losses.append(t_loss)\n    train_accs.append(t_acc)\n    \n    # Validate\n    v_loss, v_acc = validate()\n    val_losses.append(v_loss)\n    val_accs.append(v_acc)\n    \n    # Update scheduler\n    scheduler.step()\n    \n    # Check for improvement\n    if v_acc > best_val_acc:\n        best_val_acc = v_acc\n        best_val_loss = v_loss\n        no_improve_count = 0\n        torch.save(model.state_dict(), 'galaxy_efficientnet_best.pth')\n        print(f\"✅ EPOCH {epoch+1}: Train Acc: {t_acc:.2f}% | Val Acc: {v_acc:.2f}% | NEW BEST! ⭐\")\n    else:\n        no_improve_count += 1\n        print(f\"✅ EPOCH {epoch+1}: Train Acc: {t_acc:.2f}% | Val Acc: {v_acc:.2f}% | No improve ({no_improve_count}/{PATIENCE})\")\n    \n    # Early stopping\n    if no_improve_count >= PATIENCE:\n        print(f\"\\n🛑 Early stopping at epoch {epoch+1}!\")\n        break\n\ntotal_time = time.time() - start_time\nprint(\"\\n\" + \"=\"*60)\nprint(f\"🏆 Training complete!\")\nprint(f\"   Best Validation Accuracy: {best_val_acc:.2f}%\")\nprint(f\"   Total Training Time: {total_time/60:.2f} minutes\")\n\n# Save final model\ntorch.save(model.state_dict(), 'galaxy_efficientnet_final.pth')\nprint(\"\\nModels saved!\")","metadata":{"id":"cell6"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 7: Load Best Model & Evaluate","metadata":{"id":"cell7_header"}},{"cell_type":"code","source":"from sklearn.metrics import classification_report, confusion_matrix\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport numpy as np\n\n# Load best model\nmodel.load_state_dict(torch.load('galaxy_efficientnet_best.pth'))\nmodel.eval()\n\nall_preds = []\nall_labels = []\n\nprint(\"Evaluating best model...\")\n\nwith torch.no_grad():\n    for inputs, labels in val_loader:\n        inputs = inputs.to(DEVICE)\n        outputs = model(inputs)\n        _, preds = torch.max(outputs, 1)\n        all_preds.extend(preds.cpu().numpy())\n        all_labels.extend(labels.numpy())\n\nclass_names = ['Round', 'In-between', 'Cigar', 'Edge-on', 'Spiral']\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"CLASSIFICATION REPORT\")\nprint(\"=\"*60)\nprint(classification_report(all_labels, all_preds, target_names=class_names))\n\n# Confusion Matrix\ncm = confusion_matrix(all_labels, all_preds)\nplt.figure(figsize=(10, 8))\nsns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names)\nplt.xlabel('Predicted')\nplt.ylabel('True')\nplt.title('Confusion Matrix - EfficientNet-B0')\nplt.tight_layout()\nplt.savefig('confusion_matrix_efficientnet.png', dpi=150)\nplt.show()","metadata":{"id":"cell7"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 8: Test-Time Augmentation (TTA)","metadata":{"id":"cell8_header"}},{"cell_type":"code","source":"# Test-Time Augmentation for better accuracy\nprint(\"Running Test-Time Augmentation (TTA)...\")\nprint(\"This averages predictions over multiple augmented versions.\")\n\n# TTA transforms (different augmentations)\ntta_transforms = [\n    val_transform,  # Original\n    transforms.Compose([  # Rotated 90\n        transforms.Resize((224, 224)),\n        transforms.Lambda(lambda x: x.rotate(90)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ]),\n    transforms.Compose([  # Rotated 180\n        transforms.Resize((224, 224)),\n        transforms.Lambda(lambda x: x.rotate(180)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ]),\n    transforms.Compose([  # Rotated 270\n        transforms.Resize((224, 224)),\n        transforms.Lambda(lambda x: x.rotate(270)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ]),\n    transforms.Compose([  # Flipped\n        transforms.Resize((224, 224)),\n        transforms.RandomHorizontalFlip(p=1.0),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ]),\n]\n\n# Create TTA datasets\ntta_datasets = []\nfor t in tta_transforms:\n    ds = GalaxyZooDataset5Class(\n        csv_file='training_solutions_rev1.csv',\n        root_dir='images_training_rev1',\n        transform=t\n    )\n    tta_datasets.append(Subset(ds, val_indices))\n\n# Predict with TTA\nmodel.eval()\nall_probs = []\n\nfor i, ds in enumerate(tta_datasets):\n    print(f\"  Running TTA variant {i+1}/{len(tta_datasets)}...\")\n    loader = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\n    probs = []\n    \n    with torch.no_grad():\n        for inputs, _ in loader:\n            inputs = inputs.to(DEVICE)\n            outputs = torch.softmax(model(inputs), dim=1)\n            probs.append(outputs.cpu().numpy())\n    \n    all_probs.append(np.concatenate(probs))\n\n# Average probabilities\navg_probs = np.mean(all_probs, axis=0)\ntta_preds = np.argmax(avg_probs, axis=1)\n\n# TTA accuracy\ntta_acc = 100 * np.mean(tta_preds == np.array(all_labels))\nprint(f\"\\n🎯 TTA Accuracy: {tta_acc:.2f}%\")\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"CLASSIFICATION REPORT (with TTA)\")\nprint(\"=\"*60)\nprint(classification_report(all_labels, tta_preds, target_names=class_names))","metadata":{"id":"cell8"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 9: Inference Time","metadata":{"id":"cell9_header"}},{"cell_type":"code","source":"print(\"--- Inference Time ---\")\nmodel.eval()\nbatch_input = next(iter(val_loader))[0].to(DEVICE)\n\n# Warm-up\nwith torch.no_grad():\n    for _ in range(10):\n        _ = model(batch_input)\n\n# Measure\nstart_time = time.time()\nwith torch.no_grad():\n    for _ in range(100):\n        _ = model(batch_input)\nend_time = time.time()\n\ntotal_images = 100 * BATCH_SIZE\navg_time = (end_time - start_time) / total_images\nprint(f\"Inference Time per image: {avg_time * 1000:.4f} ms\")\nprint(f\"Throughput: {1/avg_time:.1f} images/second\")","metadata":{"id":"cell9"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 10: Training Curves","metadata":{"id":"cell10_header"}},{"cell_type":"code","source":"# Plot training curves\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n\nepochs_ran = len(train_losses)\n\n# Loss\nax1.plot(range(1, epochs_ran+1), train_losses, 'b-', label='Train Loss')\nax1.plot(range(1, epochs_ran+1), val_losses, 'r-', label='Val Loss')\nax1.set_xlabel('Epoch')\nax1.set_ylabel('Loss')\nax1.set_title('Training & Validation Loss')\nax1.legend()\nax1.grid(True)\n\n# Accuracy\nax2.plot(range(1, epochs_ran+1), train_accs, 'b-', label='Train Acc')\nax2.plot(range(1, epochs_ran+1), val_accs, 'r-', label='Val Acc')\nax2.axhline(y=best_val_acc, color='g', linestyle='--', label=f'Best: {best_val_acc:.2f}%')\nax2.set_xlabel('Epoch')\nax2.set_ylabel('Accuracy (%)')\nax2.set_title('Training & Validation Accuracy')\nax2.legend()\nax2.grid(True)\n\nplt.tight_layout()\nplt.savefig('training_curves_efficientnet.png', dpi=150)\nplt.show()\n\nprint(f\"\\n📊 Training ran for {epochs_ran} epochs\")\nprint(f\"   Best validation accuracy: {best_val_acc:.2f}%\")","metadata":{"id":"cell10"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 11: Parameter Statistics","metadata":{"id":"cell11_header"}},{"cell_type":"code","source":"print(\"--- Trained Parameter Statistics ---\")\nprint(\"(Selected layers)\\n\")\n\nfor name, param in model.named_parameters():\n    if 'classifier' in name or 'features.0.0' in name or 'features.8' in name:\n        print(f\"{name}: mean={param.mean().item():.6f}, std={param.std().item():.6f}\")\n\nprint(\"\\n🎉 All experiments completed!\")\nprint(\"\\nSaved files:\")\nprint(\"  - galaxy_efficientnet_best.pth\")\nprint(\"  - galaxy_efficientnet_final.pth\")\nprint(\"  - confusion_matrix_efficientnet.png\")\nprint(\"  - training_curves_efficientnet.png\")","metadata":{"id":"cell11"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Cell 12: Summary & Comparison","metadata":{"id":"cell12_header"}},{"cell_type":"code","source":"print(\"=\"*60)\nprint(\"FINAL SUMMARY\")\nprint(\"=\"*60)\nprint(f\"\")\nprint(f\"Model: EfficientNet-B0\")\nprint(f\"Parameters: ~5.3M\")\nprint(f\"Training epochs: {epochs_ran}\")\nprint(f\"Best Val Accuracy: {best_val_acc:.2f}%\")\nprint(f\"TTA Accuracy: {tta_acc:.2f}%\")\nprint(f\"Inference time: {avg_time*1000:.2f} ms/image\")\nprint(f\"Total training time: {total_time/60:.2f} minutes\")\nprint(f\"\")\nprint(\"Techniques used:\")\nprint(\"  - Transfer learning from ImageNet\")\nprint(\"  - SGD with momentum + cosine annealing\")\nprint(\"  - Class weighting for imbalanced data\")\nprint(\"  - Advanced data augmentation\")\nprint(\"  - Early stopping\")\nprint(\"  - Test-time augmentation\")\nprint(\"=\"*60)","metadata":{"id":"cell12"},"execution_count":null,"outputs":[]}]}