## Project: [Your Project Name]

### Problem Statement
What problem are you solving? Who benefits?

### Success Criteria
- Metric 1: Accuracy > 90%
- Metric 2: Inference time < 100ms
- Metric 3: Model size < 100MB

### Constraints
- Hardware: Single GPU
- Data: Public dataset + custom samples
- Timeline: 2 weeks

import os
import requests
from pathlib import Path

def download_dataset(url, save_dir):
    """Download and extract dataset."""
    save_dir = Path(save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    
    filename = url.split("/")[-1]
    filepath = save_dir / filename
    
    print(f"Downloading {url}...")
    response = requests.get(url, stream=True)
    with open(filepath, "wb") as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    
    print(f"Saved to {filepath}")
    return filepath

# Example: Download CIFAR-10
# download_dataset("https://...", "data/")

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import pandas as pd

class CustomDataset(Dataset):
    """Custom dataset for your project.
    
    Best practices embedded in this implementation:
    - CSV-based annotation for easy inspection and versioning
    - Lazy loading (images loaded on-demand, not all into memory)
    - Configurable transforms for train vs. validation
    """
    
    def __init__(self, data_dir, csv_file, transform=None):
        self.data_dir = Path(data_dir)
        self.annotations = pd.read_csv(csv_file)
        self.transform = transform or transforms.ToTensor()
    
    def __len__(self):
        return len(self.annotations)
    
    def __getitem__(self, idx):
        row = self.annotations.iloc[idx]
        
        # Load image
        image_path = self.data_dir / row["filename"]
        image = Image.open(image_path).convert("RGB")
        
        if self.transform:
            image = self.transform(image)
        
        label = row["label"]
        return image, label


def create_dataloaders(train_dir, val_dir, batch_size=32):
    """Create train and validation dataloaders.
    
    IMPORTANT: Training and validation use DIFFERENT transforms.
    - Training: random augmentation to improve generalization
    - Validation: deterministic center crop for consistent evaluation
    Never apply random augmentation to validation data -- your metrics
    would fluctuate between runs, making comparison impossible.
    """
    
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),       # Random crop forces spatial invariance
        transforms.RandomHorizontalFlip(),        # Most natural images are flip-invariant
        transforms.ColorJitter(0.2, 0.2, 0.2),   # Robustness to lighting changes
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  # ImageNet stats
    ])
    
    val_transform = transforms.Compose([
        transforms.Resize(256),        # Resize shorter edge to 256
        transforms.CenterCrop(224),    # Deterministic center crop
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])
    
    train_dataset = CustomDataset(train_dir, "train.csv", train_transform)
    val_dataset = CustomDataset(val_dir, "val.csv", val_transform)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    
    return train_loader, val_loader

import torch
import torch.nn as nn
import timm

class ProjectModel(nn.Module):
    """Your project model.
    
    Strategy: Pretrained backbone + custom classifier head.
    Phase 1: Freeze backbone, train only the head (fast convergence)
    Phase 2: Unfreeze backbone, fine-tune end-to-end (squeeze out last few %)
    """
    
    def __init__(self, num_classes, pretrained=True):
        super().__init__()
        
        # Use pretrained backbone -- timm gives you 700+ architectures
        # EfficientNet-B0 is a great starting point: strong accuracy, reasonable size
        self.backbone = timm.create_model(
            "efficientnet_b0",
            pretrained=pretrained,
            num_classes=0  # Remove classifier to use as feature extractor
        )
        
        # Custom classifier -- keep it simple to start
        # Dropout prevents overfitting when fine-tuning with limited data
        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(self.backbone.num_features, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, num_classes),
        )
    
    def forward(self, x):
        features = self.backbone(x)
        return self.classifier(features)
    
    def freeze_backbone(self):
        for param in self.backbone.parameters():
            param.requires_grad = False
    
    def unfreeze_backbone(self):
        for param in self.backbone.parameters():
            param.requires_grad = True

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.cuda.amp import autocast, GradScaler
from tqdm import tqdm
import wandb

class Trainer:
    def __init__(self, model, train_loader, val_loader, config):
        self.model = model.cuda()
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config
        
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = AdamW(model.parameters(), lr=config["lr"], weight_decay=config["weight_decay"])
        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=config["epochs"])
        self.scaler = GradScaler()
        
        self.best_val_acc = 0
        
        # Initialize wandb
        wandb.init(project="capstone", config=config)
    
    def train_epoch(self):
        self.model.train()
        total_loss = 0
        correct = 0
        total = 0
        
        pbar = tqdm(self.train_loader, desc="Training")
        for batch_idx, (images, labels) in enumerate(pbar):
            images, labels = images.cuda(), labels.cuda()
            
            self.optimizer.zero_grad()
            
            with autocast():
                outputs = self.model(images)
                loss = self.criterion(outputs, labels)
            
            self.scaler.scale(loss).backward()
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.scaler.step(self.optimizer)
            self.scaler.update()
            
            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            
            pbar.set_postfix({
                "loss": total_loss / (batch_idx + 1),
                "acc": 100. * correct / total
            })
        
        return total_loss / len(self.train_loader), 100. * correct / total
    
    @torch.no_grad()
    def validate(self):
        self.model.eval()
        total_loss = 0
        correct = 0
        total = 0
        
        for images, labels in tqdm(self.val_loader, desc="Validating"):
            images, labels = images.cuda(), labels.cuda()
            
            outputs = self.model(images)
            loss = self.criterion(outputs, labels)
            
            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
        
        return total_loss / len(self.val_loader), 100. * correct / total
    
    def train(self):
        for epoch in range(self.config["epochs"]):
            print(f"\nEpoch {epoch + 1}/{self.config['epochs']}")
            
            train_loss, train_acc = self.train_epoch()
            val_loss, val_acc = self.validate()
            self.scheduler.step()
            
            # Logging
            wandb.log({
                "train_loss": train_loss,
                "train_acc": train_acc,
                "val_loss": val_loss,
                "val_acc": val_acc,
                "lr": self.scheduler.get_last_lr()[0],
            })
            
            # Save best model
            if val_acc > self.best_val_acc:
                self.best_val_acc = val_acc
                torch.save(self.model.state_dict(), "best_model.pt")
                print(f"Saved best model with val_acc: {val_acc:.2f}%")
            
            print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%")
            print(f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%")
        
        wandb.finish()
        return self.best_val_acc


# Run training
config = {
    "epochs": 30,
    "lr": 1e-4,
    "weight_decay": 0.01,
    "batch_size": 32,
}

model = ProjectModel(num_classes=10)
trainer = Trainer(model, train_loader, val_loader, config)
trainer.train()

import numpy as np
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

def evaluate_model(model, test_loader, class_names):
    """Comprehensive model evaluation."""
    model.eval()
    all_preds = []
    all_labels = []
    all_probs = []
    
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.cuda()
            outputs = model(images)
            probs = torch.softmax(outputs, dim=1)
            _, predicted = outputs.max(1)
            
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.numpy())
            all_probs.extend(probs.cpu().numpy())
    
    # Classification report
    print("\nClassification Report:")
    print(classification_report(all_labels, all_preds, target_names=class_names))
    
    # Confusion matrix
    cm = confusion_matrix(all_labels, all_preds)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                xticklabels=class_names, yticklabels=class_names)
    plt.xlabel("Predicted")
    plt.ylabel("True")
    plt.title("Confusion Matrix")
    plt.savefig("confusion_matrix.png")
    
    return all_preds, all_labels, all_probs

from fastapi import FastAPI, File, UploadFile
from PIL import Image
import torch
import io

app = FastAPI(title="My ML Model API")

# Load model
model = ProjectModel(num_classes=10)
model.load_state_dict(torch.load("best_model.pt"))
model.eval()
model.cuda()

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # Read image
    image_bytes = await file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    
    # Preprocess
    input_tensor = transform(image).unsqueeze(0).cuda()
    
    # Inference
    with torch.no_grad():
        output = model(input_tensor)
        probs = torch.softmax(output, dim=1)
        pred_class = probs.argmax(dim=1).item()
        confidence = probs[0, pred_class].item()
    
    return {
        "class": CLASS_NAMES[pred_class],
        "confidence": f"{confidence:.2%}",
        "all_probabilities": {
            name: f"{prob:.2%}" 
            for name, prob in zip(CLASS_NAMES, probs[0].tolist())
        }
    }

@app.get("/health")
async def health():
    return {"status": "healthy"}

FROM python:3.10-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY best_model.pt .
COPY app.py .

EXPOSE 8000

CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]