Upload 5 files

Dockerfile

@@ -0,0 +1,23 @@
+# Use a lightweight Python image
+FROM python:3.9-slim
+
+WORKDIR /code
+
+# 1. Install CPU-only PyTorch (Critical for Free Tier to avoid 3GB bloat)
+RUN pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cpu
+
+# 2. Copy requirements and install dependencies
+COPY ./requirements.txt /code/requirements.txt
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# 3. Create checkpoints directory and set permissions
+# We allow any user to write here to avoid permission errors on some clouds
+RUN mkdir -p /code/checkpoints && chmod -R 777 /code/checkpoints
+
+# 4. Copy your code and the TRAINED MODEL
+COPY server.py /code/server.py
+COPY checkpoints/ /code/checkpoints/
+
+# 5. Start the Server
+# Port 7860 is the default for Hugging Face Spaces
+CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "7860"]

checkpoints/best_age_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6123e7653e7de692f0c0509b1429c5978709fc3ab0595059978e6c0d17c2cdc5
+size 45050069

main.py

@@ -0,0 +1,170 @@
+import os
+import glob
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader, random_split
+from torchvision import models, transforms as T
+from PIL import Image
+from tqdm import tqdm # Progress bar
+
+# =================CONFIGURATION=================
+# Update this path based on your screenshot
+DATA_DIR = "dataset/UTKFace" 
+CHECKPOINT_DIR = "checkpoints"
+BATCH_SIZE = 64
+EPOCHS = 10
+LEARNING_RATE = 1e-4
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# ===============================================
+
+# We must use a tiny class for Data because PyTorch requires it
+# But we keep it minimal (just loading logic)
+class SimpleFaceDataset(Dataset):
+    def __init__(self, root_dir, transform=None):
+        self.transform = transform
+        # Find all files ending in .jpg (some in your screenshot have .chip.jpg)
+        self.files = glob.glob(os.path.join(root_dir, "*.jpg*"))
+        print(f"Dataset: Found {len(self.files)} images.")
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, idx):
+        filepath = self.files[idx]
+        filename = os.path.basename(filepath)
+        
+        # 1. Parse Age from Filename
+        # Example: "26_0_1_2017.jpg" -> "26" is the first part
+        try:
+            age = int(filename.split('_')[0])
+        except:
+            age = 25 # Default safe fallback
+            
+        # 2. Load Image
+        image = Image.open(filepath).convert('RGB')
+        
+        if self.transform:
+            image = self.transform(image)
+            
+        # Return image and age as a Float (needed for regression)
+        return image, torch.tensor(age, dtype=torch.float32)
+
+def get_model():
+    """Loads ResNet18 and changes output to 1 number (Age)."""
+    print("Loading ResNet18...")
+    model = models.resnet18(weights="DEFAULT")
+    
+    # Change the final layer
+    # Old: Output 1000 classes (ImageNet)
+    # New: Output 1 number (Age)
+    num_features = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Linear(num_features, 128),
+        nn.ReLU(),
+        nn.Dropout(0.2), # Helps prevent memorization
+        nn.Linear(128, 1)
+    )
+    return model.to(DEVICE)
+
+def train_one_epoch(model, loader, criterion, optimizer):
+    model.train()
+    running_loss = 0.0
+    total_error = 0.0 # To track how many years we are off by
+    count = 0
+    
+    loop = tqdm(loader, desc="Training", leave=False)
+    for images, ages in loop:
+        images, ages = images.to(DEVICE), ages.to(DEVICE)
+        
+        # Forward
+        predictions = model(images).squeeze() # remove extra dim
+        loss = criterion(predictions, ages)
+        
+        # Backward
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        
+        # Stats
+        running_loss += loss.item()
+        # Calculate Absolute Error (e.g., predicted 25, real 30 -> error 5)
+        total_error += torch.abs(predictions - ages).sum().item()
+        count += ages.size(0)
+        
+        loop.set_description(f"Loss: {loss.item():.2f}")
+
+    avg_loss = running_loss / len(loader)
+    avg_mae = total_error / count # Mean Absolute Error
+    return avg_loss, avg_mae
+
+def evaluate(model, loader, criterion):
+    model.eval()
+    running_loss = 0.0
+    total_error = 0.0
+    count = 0
+    
+    with torch.no_grad():
+        for images, ages in loader:
+            images, ages = images.to(DEVICE), ages.to(DEVICE)
+            
+            predictions = model(images).squeeze()
+            loss = criterion(predictions, ages)
+            
+            running_loss += loss.item()
+            total_error += torch.abs(predictions - ages).sum().item()
+            count += ages.size(0)
+            
+    avg_loss = running_loss / len(loader)
+    avg_mae = total_error / count
+    return avg_loss, avg_mae
+
+# ================= MAIN SCRIPT =================
+if __name__ == "__main__":
+    # 1. Setup Transforms (Resize to 200x200 standard)
+    transform = T.Compose([
+        T.Resize((200, 200)),
+        T.ToTensor(),
+        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    
+    # 2. Setup Data
+    full_dataset = SimpleFaceDataset(DATA_DIR, transform=transform)
+    
+    # 80/20 Split
+    train_size = int(0.8 * len(full_dataset))
+    val_size = len(full_dataset) - train_size
+    train_ds, val_ds = random_split(full_dataset, [train_size, val_size])
+    
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
+    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, num_workers=4)
+    
+    # 3. Setup Model, Loss, Optimizer
+    model = get_model()
+    # MSELoss is standard for predicting numbers (Regression)
+    criterion = nn.MSELoss() 
+    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
+    
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    
+    print(f"Starting training on {DEVICE}...")
+    
+    best_mae = float('inf') # Lower is better
+    
+    for epoch in range(EPOCHS):
+        print(f"\nEpoch {epoch+1}/{EPOCHS}")
+        
+        train_loss, train_mae = train_one_epoch(model, train_loader, criterion, optimizer)
+        print(f"Train Loss: {train_loss:.2f} | Avg Error: {train_mae:.1f} years")
+        
+        val_loss, val_mae = evaluate(model, val_loader, criterion)
+        print(f"Val Loss:   {val_loss:.2f} | Avg Error: {val_mae:.1f} years")
+        
+        # Save Best
+        if val_mae < best_mae:
+            best_mae = val_mae
+            save_path = os.path.join(CHECKPOINT_DIR, "best_age_model.pth")
+            torch.save(model.state_dict(), save_path)
+            print(f"🔥 Model saved! (Off by only {val_mae:.1f} years on avg)")
+
+    print("Done!")

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+python-multipart
+torch
+torchvision
+pillow
+tqdm

server.py

@@ -0,0 +1,89 @@
+import io
+import os
+import glob
+import torch
+import torch.nn as nn
+from torchvision import models, transforms as T
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File, HTTPException
+import uvicorn
+
+app = FastAPI()
+
+# =================CONFIGURATION=================
+CHECKPOINT_PATH = "checkpoints/best_age_model.pth"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# ===============================================
+
+# 1. Define the Model Architecture (Must match training EXACTLY)
+def get_model():
+    # Load standard ResNet18 structure
+    model = models.resnet18(weights=None) 
+    
+    # Reconstruct the final layer exactly as we defined in training
+    num_features = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Linear(num_features, 128),
+        nn.ReLU(),
+        nn.Dropout(0.2),
+        nn.Linear(128, 1)
+    )
+    return model
+
+# 2. Load Weights
+print(f"Initializing model on {DEVICE}...")
+model = get_model()
+
+if os.path.exists(CHECKPOINT_PATH):
+    print(f"Loading weights from {CHECKPOINT_PATH}...")
+    # Load the state dictionary
+    state_dict = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
+    model.load_state_dict(state_dict)
+    model.to(DEVICE)
+    model.eval() # Important: Turn off Dropout for inference
+    print("✅ Model loaded successfully!")
+else:
+    print(f"⚠️ Warning: Checkpoint not found at {CHECKPOINT_PATH}")
+    print("Inference will use random weights (garbage output).")
+
+# 3. Define Transforms (Must match training)
+transform = T.Compose([
+    T.Resize((200, 200)),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+@app.get("/")
+def home():
+    return {"status": "Age Predictor API is running"}
+
+@app.post("/predict")
+async def predict_age(file: UploadFile = File(...)):
+    """
+    Accepts an image file and returns predicted age.
+    """
+    try:
+        # 1. Read and Process Image
+        content = await file.read()
+        image = Image.open(io.BytesIO(content)).convert("RGB")
+        
+        # 2. Transform
+        input_tensor = transform(image).unsqueeze(0).to(DEVICE)
+        
+        # 3. Inference
+        with torch.no_grad():
+            prediction = model(input_tensor).squeeze()
+            predicted_age = prediction.item()
+            
+        # 4. Return Result
+        return {
+            "filename": file.filename,
+            "predicted_age": round(predicted_age, 1) # Round to 1 decimal
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    # Host 0.0.0.0 is needed for cloud environments/Docker
+    uvicorn.run(app, host="0.0.0.0", port=8000)