Upload 3 files

Dockerfile

@@ -0,0 +1,24 @@
+# Use an official Python runtime
+FROM python:3.10-slim
+
+# Set up a new user named "user" with user ID 1000
+# (Required for Hugging Face Spaces to prevent permission errors)
+RUN useradd -m -u 1000 user
+USER user
+
+# Set environmental variables
+ENV PATH="/home/user/.local/bin:$PATH"
+WORKDIR /home/user/app
+
+# Copy requirements and install
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy your backend code
+COPY --chown=user . .
+
+# Hugging Face Spaces route web traffic to port 7860
+EXPOSE 7860
+
+# Start the FastAPI app on port 7860
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,327 @@
+import os
+import uuid
+import pickle
+from typing import List
+
+try:
+    from fastapi import FastAPI, File, Form, UploadFile, HTTPException
+    from fastapi.middleware.cors import CORSMiddleware
+except ImportError as exc:
+    raise ImportError(
+        "FastAPI is required to run this application. Install it with 'pip install fastapi'."
+    ) from exc
+
+import torch
+import torchvision.models as models
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import train_test_split  # Added for accuracy scoring
+import io
+
+# ── App Initialization ───────────────────────────────────────────────────────
+app = FastAPI(title="Teachable Machine Backend")
+
+# ── CORS Configuration ───────────────────────────────────────────────────────
+# Enables file uploads and API calls from frontend (running on different origin/port)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allow requests from any origin
+    allow_credentials=True,
+    allow_methods=["*"],  # Allow all HTTP methods
+    allow_headers=["*"],  # Allow all headers
+)
+
+DATASET_DIR = os.path.join(os.path.dirname(__file__), "dataset")
+MODEL_PATH  = os.path.join(os.path.dirname(__file__), "model.pkl")
+
+
+# ── Shared ML Setup (runs once at startup) ───────────────────────────────────
+# Loading the model once here means every request reuses the same object in
+# memory instead of reloading it from disk each time — much faster.
+device = torch.device("cpu")
+
+backbone = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.DEFAULT)
+# Remove the final classifier layer — we only want feature extraction.
+# The 960 numbers it outputs describe the image content without predicting a category.
+backbone.classifier = torch.nn.Identity()
+backbone.eval()  # Disables dropout — we are inferring, not training the backbone
+
+# These values MUST be identical during training and prediction.
+# 224x224 = the size MobileNetV3 was designed for.
+# mean/std = ImageNet dataset statistics the model was pre-trained on.
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+
+# ── Helper: Extract features from a PIL image ────────────────────────────────
+def extract_features(pil_image: Image.Image) -> np.ndarray:
+    """
+    Passes an image through MobileNetV3 and returns a 960-number feature vector.
+    Used by both /train and /predict to guarantee identical preprocessing.
+    """
+    image = pil_image.convert("RGB")   # Handles RGBA/grayscale images safely
+    tensor = transform(image)
+    tensor = tensor.unsqueeze(0)       # (3,224,224) → (1,3,224,224) — adds batch dim
+
+    with torch.no_grad():              # No gradients needed — saves memory & time
+        features = backbone(tensor)
+
+    return features.squeeze().numpy() # (1,960) → (960,) numpy array for sklearn
+
+
+# ── Health Check ─────────────────────────────────────────────────────────────
+@app.get("/")
+def health_check():
+    return {"status": "Backend is running!"}
+
+
+# ── Milestone 1: Upload images ───────────────────────────────────────────────
+@app.post("/upload-sample")
+async def upload_sample(
+    class_name: str = Form(...),
+    files: List[UploadFile] = File(...)
+):
+    """
+    Accepts a class label + a batch of images.
+    Saves each image into dataset/<class_name>/ with a random UUID filename.
+    """
+    class_name = class_name.strip().replace(" ", "_")
+
+    if not class_name:
+        raise HTTPException(status_code=400, detail="class_name cannot be empty.")
+
+    class_folder = os.path.join(DATASET_DIR, class_name)
+    os.makedirs(class_folder, exist_ok=True)
+
+    if not files:
+        raise HTTPException(status_code=400, detail="At least one image file is required.")
+
+    saved_files = []
+
+    for file in files:
+        if not file.content_type.startswith("image/"):
+            raise HTTPException(
+                status_code=400,
+                detail=f"File '{file.filename}' is not an image. Only image files are accepted."
+            )
+
+        extension = os.path.splitext(file.filename)[1] or ".jpg"
+        random_filename = f"{uuid.uuid4()}{extension}"
+        save_path = os.path.join(class_folder, random_filename)
+
+        contents = await file.read()
+        with open(save_path, "wb") as f:
+            f.write(contents)
+
+        saved_files.append(random_filename)
+
+    return {
+        "message": f"Uploaded {len(saved_files)} image(s) to class '{class_name}'",
+        "class": class_name,
+        "saved_files": saved_files
+    }
+
+
+# ── Milestone 1 Bonus: Dataset info ─────────────────────────────────────────
+@app.get("/dataset-info")
+def dataset_info():
+    if not os.path.exists(DATASET_DIR):
+        return {"classes": {}, "total_images": 0}
+
+    summary = {}
+    for class_name in os.listdir(DATASET_DIR):
+        class_path = os.path.join(DATASET_DIR, class_name)
+        if os.path.isdir(class_path):
+            summary[class_name] = len(os.listdir(class_path))
+
+    return {
+        "classes": summary,
+        "total_images": sum(summary.values())
+    }
+
+
+# ── Milestone 2: Train model ─────────────────────────────────────────────────
+@app.post("/train")
+def train_model():
+    """
+    Scans dataset/, extracts MobileNetV3 features from every image,
+    trains a LogisticRegression classifier, and saves it to model.pkl.
+    """
+
+    # ── Step 1: Validate dataset exists ──────────────────────────────────────
+    if not os.path.exists(DATASET_DIR):
+        raise HTTPException(
+            status_code=400,
+            detail="No dataset found. Please upload images first."
+        )
+
+    classes = [
+        d for d in os.listdir(DATASET_DIR)
+        if os.path.isdir(os.path.join(DATASET_DIR, d))
+    ]
+
+    # Classifier needs at least 2 classes — it learns to DISTINGUISH between them.
+    # With only 1 class there is nothing to distinguish.
+    if len(classes) < 2:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Need at least 2 classes to train. You currently have: {classes}"
+        )
+
+    X = []  # Feature vectors — one row per image
+    y = []  # Labels         — one entry per image, matched by index to X
+
+    # ── Step 2: Extract features from every image ────────────────────────────
+    for class_name in classes:
+        class_folder = os.path.join(DATASET_DIR, class_name)
+        image_files  = os.listdir(class_folder)
+
+        if len(image_files) == 0:
+            continue  # Skip empty class folders silently
+
+        for filename in image_files:
+            image_path = os.path.join(class_folder, filename)
+            try:
+                img      = Image.open(image_path)
+                features = extract_features(img)
+                X.append(features)
+                y.append(class_name)
+            except Exception as e:
+                # One corrupted image should not kill the whole training run
+                print(f"Skipping {filename}: {e}")
+                continue
+
+    if len(X) == 0:
+        raise HTTPException(
+            status_code=400,
+            detail="No valid images found in dataset."
+        )
+
+    X = np.array(X)  # Shape: (num_images, 960)
+    y = np.array(y)  # Shape: (num_images,)
+
+    # ── Step 3: Train the classifier ─────────────────────────────────────────
+    
+    # NEW: Split the data to calculate a real accuracy metric.
+    # We added a safety net: if there are fewer than 5 images total, we test 
+    # on the training data so it doesn't crash during a live presentation.
+    if len(X) >= 5:
+        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+    else:
+        X_train, X_test, y_train, y_test = X, X, y, y
+
+    # Why LogisticRegression?
+    # MobileNetV3 already converted images into meaningful 960-number vectors.
+    # LogisticRegression just finds the boundary between those vectors.
+    # It trains in under a second, works with very few images, and needs no GPU.
+    # max_iter=1000 prevents ConvergenceWarning on small datasets.
+    classifier = LogisticRegression(max_iter=1000)
+    classifier.fit(X_train, y_train)
+    
+    # Calculate overall accuracy
+    accuracy = classifier.score(X_test, y_test)
+
+    # ── Step 4: Save classifier + class list to disk ─────────────────────────
+    # We save classes explicitly so the /predict endpoint can map
+    # numeric outputs back to human-readable label names.
+    model_data = {
+        "classifier": classifier,
+        "classes": classes
+    }
+    with open(MODEL_PATH, "wb") as f:
+        pickle.dump(model_data, f)
+
+    return {
+        "message": "Training complete!",
+        "classes": classes,
+        "total_images": len(X),
+        "accuracy": round(accuracy * 100, 2),  # Returned safely to the frontend!
+        "model_saved_at": MODEL_PATH
+    }
+
+
+# ── Milestone 3: Predict endpoint ────────────────────────────────────────────
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    """
+    Accepts a single image, runs it through MobileNetV3 + the trained
+    LogisticRegression classifier, and returns the predicted class
+    with a confidence score for every class.
+    """
+
+    # ── Step 1: Check model exists ────────────────────────────────────────────
+    # If the user hits /predict before ever running /train, model.pkl won't
+    # exist yet. We catch this early with a clear message instead of a crash.
+    if not os.path.exists(MODEL_PATH):
+        raise HTTPException(
+            status_code=400,
+            detail="No trained model found. Please call /train first."
+        )
+
+    # ── Step 2: Validate the uploaded file is an image ────────────────────────
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(
+            status_code=400,
+            detail=f"File '{file.filename}' is not an image. Only image files are accepted."
+        )
+
+    # ── Step 3: Load the saved model from disk ────────────────────────────────
+    # We reload model.pkl on every prediction request.
+    # Why not load it once at startup like the backbone?
+    # Because model.pkl gets replaced every time /train is called.
+    # If we cached it at startup, predictions would use the OLD model
+    # even after the user retrains — a subtle but serious bug.
+    with open(MODEL_PATH, "rb") as f:
+        model_data = pickle.load(f)
+
+    classifier = model_data["classifier"]
+    classes    = model_data["classes"]
+
+    # ── Step 4: Read and decode the uploaded image ────────────────────────────
+    # file.read() gives us raw bytes. We wrap them in BytesIO so PIL
+    # can treat the bytes like a file on disk — no temp file needed.
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents))
+
+    # ── Step 5: Extract features using the SAME function used during training ─
+    # This is the most important consistency rule in the whole project.
+    # If training used 224x224 + ImageNet normalization, prediction MUST too.
+    # extract_features() guarantees this since both phases call the same code.
+    features = extract_features(image)
+
+    # ── Step 6: Run prediction ────────────────────────────────────────────────
+    # features is shape (960,) — we reshape to (1, 960) because sklearn
+    # expects a 2D array: (number_of_samples, number_of_features)
+    features_2d = features.reshape(1, -1)
+
+    # predict() returns the winning class label e.g. ["cat"]
+    predicted_class = classifier.predict(features_2d)[0]
+
+    # predict_proba() returns confidence scores for ALL classes e.g. [0.82, 0.18]
+    # Each number = how confident the model is that this image belongs to that class.
+    # They always sum to 1.0 (100%).
+    probabilities = classifier.predict_proba(features_2d)[0]
+
+    # ── Step 7: Build a clean confidence scores dict ──────────────────────────
+    # zip(classes, probabilities) pairs each class name with its score:
+    # e.g. {"cat": 0.82, "dog": 0.18}
+    # round(..., 4) keeps it readable: 0.8173 instead of 0.81734521938...
+    # float() converts numpy float32 → Python float so JSON can serialize it
+    confidence_scores = {
+        cls: round(float(prob), 4)
+        for cls, prob in zip(classifier.classes_, probabilities)
+    }
+
+    return {
+        "predicted_class": predicted_class,
+        "confidence": round(float(max(probabilities)), 4),
+        "all_scores": confidence_scores
+    }

requirements.txt

@@ -0,0 +1,9 @@
+fastapi
+uvicorn[standard]
+python-multipart
+torch
+torchvision
+scikit-learn
+Pillow
+streamlit
+requests