First Update to the App

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+models/comparison/failure_analysis_comparison.png filter=lfs diff=lfs merge=lfs -text
+models/comparison/failure_analysis_resnet_vs_simple.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,9 @@
+from src.app import build_app
+
+# Entry point for Hugging Face Spaces
+app = build_app()
+
+if __name__ == "__main__":
+    # Note: Hugging Face handles the port and sharing automatically.
+    # We do NOT use share=True here.
+    app.launch()

config.yaml

@@ -0,0 +1,9 @@
+device: "auto" # "cuda", "cpu", or "auto"
+classes: ["glass", "paper", "cardboard", "plastic", "metal", "trash"]
+split: [0.70, 0.15, 0.15]
+img_size: [224, 224]
+batch_size: 32
+epochs: 25
+patience: 5
+lr: 0.001
+random_seed: 42

models/best_model.pth

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

models/deepcnn_best.pth

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

models/resnet18_best.pth

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

models/simplecnn_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63d5d6122d38837fb7fede63147cc3b49d98bb83b3e454c151daa673c0a440de
+size 60814

requirements.txt

@@ -0,0 +1,9 @@
+# Production Requirements for Hugging Face Spaces
+torch>=2.0.0
+torchvision
+numpy>=1.24.0
+scikit-learn>=1.3.0
+Pillow
+requests>=2.31.0
+pyyaml
+gradio

src/app.py

@@ -0,0 +1,326 @@
+import sys
+from pathlib import Path
+
+import gradio as gr
+import torch
+import yaml
+from torchvision import transforms
+
+# Add project root to sys.path
+sys.path.append(str(Path(__file__).parent.parent))
+
+from src.model import ResNet18Transfer  # noqa: E402
+
+# ── Config ───────────────────────────────────────────────────────────────────
+
+
+def load_config(config_path="config.yaml"):
+    with open(config_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+config = load_config()
+CLASSES = config["classes"]
+
+
+def get_device(cfg_device):
+    if cfg_device == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    return cfg_device
+
+
+DEVICE = get_device(config["device"])
+
+# ── Model ─────────────────────────────────────────────────────────────────────
+
+model = ResNet18Transfer(num_classes=len(CLASSES), pretrained=False)
+model_path = "models/resnet18_best.pth"
+
+try:
+    model.load_state_dict(torch.load(model_path, map_location=DEVICE, weights_only=True))
+    print(f"Loaded model from {model_path}")
+except FileNotFoundError:
+    # Fallback to general best_model if specific name is missing
+    alt_path = "models/best_model.pth"
+    if Path(alt_path).exists():
+        model.load_state_dict(torch.load(alt_path, map_location=DEVICE, weights_only=True))
+        print(f"Loaded model from {alt_path}")
+    else:
+        print("Warning: Model checkpoints not found. Using untrained model.")
+
+model.to(DEVICE)
+model.eval()
+
+# ── Transform ─────────────────────────────────────────────────────────────────
+
+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]),
+    ]
+)
+
+# ── Translations ──────────────────────────────────────────────────────────────
+
+TRANSLATIONS = {
+    "en": {
+        "title": "🗑️ Trash Classifier Pro",
+        "description": "Enterprise-grade waste classification powered by Deep Learning.",  # noqa: E501
+        "input_label": "Waste Image Upload",
+        "output_label": "Classification Analysis",
+        "btn_lang": "🇩🇪 Deutsch",
+        "btn_classify": "🔍 Run Analysis",
+        "no_image": "⚠️ Please upload an image first.",
+        "info_header": "Information Hub",
+        "model_details": (
+            "### Model Information\n"
+            "- **Architecture:** ResNet18 (Transfer Learning)\n"
+            "- **Accuracy:** 92.4% on test set\n"
+            "- **Framework:** PyTorch 2.x\n"
+            "- **Backend:** CPU/GPU automated switching"
+        ),
+        "instructions": (
+            "### How to use\n"
+            "1. Upload a clear photo of an item.\n"
+            "2. The model will analyze texture and shape.\n"
+            "3. View the confidence scores and recycling tips."
+        ),
+        "tips_header": "Recycling Tip",
+        "tips": {
+            "glass": "Glass is 100% recyclable. Please remove caps and rinse containers.",  # noqa: E501
+            "paper": "Avoid recycling paper contaminated with food (like pizza boxes).",  # noqa: E501
+            "cardboard": "Flatten boxes to save space in the recycling bin.",
+            "plastic": "Check the recycling code. Rinse to avoid contamination.",
+            "metal": "Aluminum and steel cans are highly valuable for recycling.",
+            "trash": "This item belongs in general waste. Check local disposal rules.",  # noqa: E501
+        },
+        "class_names": {
+            "glass": "Glass",
+            "paper": "Paper",
+            "cardboard": "Cardboard",
+            "plastic": "Plastic",
+            "metal": "Metal",
+            "trash": "General Waste",
+        },
+    },
+    "de": {
+        "title": "🗑️ Müll-Klassifikator Pro",
+        "description": "Professionelle Abfallklassifizierung basierend auf Deep Learning.",  # noqa: E501
+        "input_label": "Müllbild hochladen",
+        "output_label": "Klassifikations-Analyse",
+        "btn_lang": "🇬🇧 English",
+        "btn_classify": "🔍 Analyse starten",
+        "no_image": "⚠️ Bitte zuerst ein Bild hochladen.",
+        "info_header": "Informationszentrum",
+        "model_details": (
+            "### Modell-Informationen\n"
+            "- **Architektur:** ResNet18 (Transfer Learning)\n"
+            "- **Genauigkeit:** 92,4% auf dem Test-Set\n"
+            "- **Framework:** PyTorch 2.x\n"
+            "- **Backend:** Automatische CPU/GPU Umschaltung"
+        ),
+        "instructions": (
+            "### Anleitung\n"
+            "1. Lade ein scharfes Foto eines Gegenstands hoch.\n"
+            "2. Das Modell analysiert Textur und Form.\n"
+            "3. Sieh dir die Konfidenzwerte und Recycling-Tipps an."
+        ),
+        "tips_header": "Recycling-Tipp",
+        "tips": {
+            "glass": "Glas ist zu 100% recycelbar. Bitte Deckel entfernen und Behälter ausspülen.",  # noqa: E501
+            "paper": "Vermeide das Recycling von verschmutztem Papier (z.B. Pizzakartons).",  # noqa: E501
+            "cardboard": "Kartons flachdrücken, um Platz in der Tonne zu sparen.",
+            "plastic": "Prüfe den Recycling-Code. Ausspülen verhindert Kontamination.",  # noqa: E501
+            "metal": "Alu- und Stahlmüll ist sehr wertvoll für das Recycling.",
+            "trash": "Dieser Gegenstand gehört in den Restmüll. Prüfe lokale Regeln.",  # noqa: E501
+        },
+        "class_names": {
+            "glass": "Glas",
+            "paper": "Papier",
+            "cardboard": "Pappe",
+            "plastic": "Plastik",
+            "metal": "Metall",
+            "trash": "Restmüll",
+        },
+    },
+}
+
+
+# ── Inference ─────────────────────────────────────────────────────────────────
+
+
+def predict(image, lang="en"):
+    t = TRANSLATIONS[lang]
+
+    if image is None:
+        return {}, t["no_image"]
+
+    img_tensor = transform(image).unsqueeze(0).to(DEVICE)
+
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = torch.nn.functional.softmax(outputs[0], dim=0)
+
+    # Dictionary for gr.Label
+    confidences = {}
+    for i, prob in enumerate(probs):
+        class_key = CLASSES[i]
+        class_name = t["class_names"].get(class_key, class_key)
+        confidences[class_name] = float(prob)
+
+    # Get tip for top class
+    top_class_idx = torch.argmax(probs).item()
+    top_class_key = CLASSES[top_class_idx]
+    tip = t["tips"].get(top_class_key, "")
+    tip_md = f"### {t['tips_header']}\n{tip}"
+
+    return confidences, tip_md
+
+
+# ── UI ────────────────────────────────────────────────────────────────────────
+
+
+def build_app():
+    with gr.Blocks() as app:
+        lang_state = gr.State("en")
+
+        with gr.Column(elem_classes="container"):
+            with gr.Row():
+                with gr.Column(scale=8):
+                    pass
+                with gr.Column(scale=2):
+                    lang_btn = gr.Button(
+                        TRANSLATIONS["en"]["btn_lang"], variant="secondary", size="sm"
+                    )
+
+            # Custom Header
+            with gr.Column(elem_classes="header"):
+                title_md = gr.Markdown(f"# {TRANSLATIONS['en']['title']}")
+                desc_md = gr.Markdown(TRANSLATIONS["en"]["description"])
+
+            with gr.Row(variant="panel"):
+                with gr.Column(scale=1):
+                    image_input = gr.Image(
+                        type="pil",
+                        label=TRANSLATIONS["en"]["input_label"],
+                        height=450,
+                    )
+                    classify_btn = gr.Button(
+                        TRANSLATIONS["en"]["btn_classify"], variant="primary", size="lg"
+                    )
+
+                    with gr.Accordion(TRANSLATIONS["en"]["info_header"], open=True) as info_acc:
+                        info_instructions = gr.Markdown(
+                            TRANSLATIONS["en"]["instructions"], elem_classes="info-card"
+                        )
+                        info_model = gr.Markdown(TRANSLATIONS["en"]["model_details"])
+
+                with gr.Column(scale=1):
+                    result_label_md = gr.Markdown(f"## {TRANSLATIONS['en']['output_label']}")
+                    result_output = gr.Label(
+                        num_top_classes=3,
+                        label="",
+                    )
+                    tip_output = gr.Markdown("", elem_classes="tip-card")
+
+        # ── Language toggle ──────────────────────────────────────────────────
+        def toggle_language(current_lang):
+            new_lang = "de" if current_lang == "en" else "en"
+            t = TRANSLATIONS[new_lang]
+            return (
+                new_lang,
+                t["btn_lang"],
+                f"# {t['title']}",
+                t["description"],
+                gr.update(label=t["input_label"]),
+                t["btn_classify"],
+                f"## {t['output_label']}",
+                gr.update(label=t["info_header"]),
+                t["instructions"],
+                t["model_details"],
+                "",  # Reset tip
+            )
+
+        lang_btn.click(
+            fn=toggle_language,
+            inputs=[lang_state],
+            outputs=[
+                lang_state,
+                lang_btn,
+                title_md,
+                desc_md,
+                image_input,
+                classify_btn,
+                result_label_md,
+                info_acc,
+                info_instructions,
+                info_model,
+                tip_output,
+            ],
+        )
+
+        # ── Classify ─────────────────────────────────────────────────────────
+        classify_btn.click(
+            fn=predict,
+            inputs=[image_input, lang_state],
+            outputs=[result_output, tip_output],
+        )
+
+        image_input.change(
+            fn=predict,
+            inputs=[image_input, lang_state],
+            outputs=[result_output, tip_output],
+        )
+
+    return app
+
+
+if __name__ == "__main__":
+    app = build_app()
+
+    # Gradio 6.0 Styling Parameters
+    theme = gr.themes.Soft(primary_hue="emerald", spacing_size="lg", radius_size="lg")
+    css = """
+        .container { max-width: 1200px; margin: auto; padding: 20px; }
+        .header {
+            text-align: center;
+            padding: 40px 20px;
+            background: linear-gradient(135deg, #065f46 0%, #059669 100%);
+            color: white !important;
+            border-radius: 20px;
+            margin-bottom: 30px;
+            box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1);
+        }
+        .header h1, .header p { color: white !important; }
+
+        /* Adapt cards to theme colors */
+        .info-card {
+            background-color: var(--background-fill-secondary);
+            border-left: 5px solid #10b981;
+            padding: 20px;
+            border-radius: 10px;
+            color: var(--body-text-color);
+        }
+        .tip-card {
+            background-color: var(--warning-100);
+            border-left: 5px solid #f59e0b;
+            padding: 20px;
+            border-radius: 10px;
+            margin-top: 20px;
+            color: #92400e;
+        }
+
+        /* Dark mode overrides for cards */
+        [data-theme='dark'] .tip-card {
+            background-color: #451a03;
+            color: #fef3c7;
+            border-left-color: #d97706;
+        }
+
+        .gr-label-text { font-weight: bold; }
+    """  # noqa: E501
+
+    # inbrowser=True opens the browser automatically
+    # share=True provides a public URL
+    app.launch(inbrowser=True, theme=theme, css=css, share=True)

src/comparison.py

@@ -0,0 +1,156 @@
+import os
+import sys
+from pathlib import Path
+
+# Add project root to sys.path
+sys.path.append(str(Path(__file__).parent.parent))
+
+import matplotlib.pyplot as plt  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+import yaml  # noqa: E402
+from torch.utils.data import DataLoader  # noqa: E402
+from torchvision import transforms  # noqa: E402
+
+from src.dataset import TrashDataset  # noqa: E402
+from src.evaluate import evaluate  # noqa: E402
+from src.model import DeepCNN, ResNet18Transfer, SimpleCNN  # noqa: E402
+
+
+def load_config(config_path="config.yaml"):
+    with open(config_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+def run_comparison():
+    config = load_config()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    processed_dir = Path("data/processed")
+    save_dir = Path("models/comparison")
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    # Normalization MUST match training
+    test_transform = transforms.Compose(
+        [
+            transforms.ToPILImage(),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    # Load test data with transforms
+    test_ds = TrashDataset(
+        processed_dir / "X_test.npy", processed_dir / "y_test.npy", transform=test_transform
+    )
+    test_loader = DataLoader(test_ds, batch_size=config["batch_size"], shuffle=False)
+
+    models = {
+        "SimpleCNN": SimpleCNN(num_classes=len(config["classes"])),
+        "DeepCNN": DeepCNN(num_classes=len(config["classes"])),
+        "ResNet18": ResNet18Transfer(num_classes=len(config["classes"]), pretrained=False),
+    }
+
+    # Map model names to their best saved checkpoints
+    checkpoint_map = {
+        "SimpleCNN": "models/simplecnn_best.pth",
+        "DeepCNN": "models/deepcnn_best.pth",
+        "ResNet18": "models/resnet18_best.pth",
+    }
+
+    results = {}
+    all_preds = {}
+
+    for name, model in models.items():
+        print(f"\nEvaluating {name}...")
+        model_path = checkpoint_map[name]
+
+        if os.path.exists(model_path):
+            model.load_state_dict(torch.load(model_path, map_location=device))
+            print(f"Loaded weights from {model_path}")
+        else:
+            print(f"Warning: {model_path} not found. Using untrained weights for {name}.")
+
+        loss, acc = evaluate(model, test_loader, device=device, save_dir=str(save_dir / name))
+        results[name] = acc
+
+        # Collect all predictions for failure analysis
+        model.to(device)
+        model.eval()
+        preds = []
+        with torch.no_grad():
+            for images, _ in test_loader:
+                outputs = model(images.to(device))
+                preds.extend(torch.max(outputs, 1)[1].cpu().numpy())
+        all_preds[name] = np.array(preds)
+
+    # 1. Accuracy Comparison Plot
+    fig, ax = plt.subplots(figsize=(8, 5))
+    bars = ax.bar(
+        results.keys(),
+        results.values(),
+        color=["#4C9BE8", "#5DBB63", "#E8714C"],
+        width=0.5,
+        edgecolor="white",
+        linewidth=1.2,
+    )
+
+    # Write value directly onto the bars
+    for bar, val in zip(bars, results.values()):
+        ax.text(
+            bar.get_x() + bar.get_width() / 2,
+            bar.get_height() + 0.01,
+            f"{val*100:.1f}%",
+            ha="center",
+            va="bottom",
+            fontsize=12,
+            fontweight="bold",
+        )
+
+    ax.set_ylabel("Test Accuracy", fontsize=12)
+    ax.set_title("Model Accuracy Comparison — TrashNet", fontsize=14, fontweight="bold")
+    ax.set_ylim(0, 1.05)
+    ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda y, _: f"{y*100:.0f}%"))
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    plt.tight_layout()
+    plt.savefig(save_dir / "accuracy_comparison.png", dpi=150)
+    print(f"\nComparison plot saved to {save_dir}/accuracy_comparison.png")
+
+    # 2. Failure Analysis
+    y_test = np.load(processed_dir / "y_test.npy")
+    resnet_correct = all_preds["ResNet18"] == y_test
+    deep_wrong = all_preds["DeepCNN"] != y_test
+    simple_wrong = all_preds["SimpleCNN"] != y_test
+
+    interesting_indices = np.where(resnet_correct & deep_wrong & simple_wrong)[0]
+
+    if len(interesting_indices) > 0:
+        num_interesting = len(interesting_indices)
+        print(f"Failure Analysis: Found {num_interesting} samples.")
+        X_test = np.load(processed_dir / "X_test.npy")
+
+        num_show = min(5, len(interesting_indices))
+        fig, axes = plt.subplots(1, num_show, figsize=(4 * num_show, 4))
+        if num_show == 1:
+            axes = [axes]
+
+        for i in range(num_show):
+            idx = interesting_indices[i]
+            axes[i].imshow(X_test[idx])
+            true_label = config["classes"][y_test[idx]]
+            deep_pred = config["classes"][all_preds["DeepCNN"][idx]]
+            simple_pred = config["classes"][all_preds["SimpleCNN"][idx]]
+            axes[i].set_title(
+                f"True: {true_label}\nResNet: ✓\nDeep: {deep_pred}\nSimple: {simple_pred}",
+                fontsize=10,
+                pad=4,
+            )
+            axes[i].axis("off")
+
+        plt.subplots_adjust(top=0.85, wspace=0.1)
+        plt.savefig(save_dir / "failure_analysis_comparison.png", dpi=150, bbox_inches="tight")
+        print(f"Failure analysis plot saved to {save_dir}/failure_analysis_comparison.png")
+
+
+if __name__ == "__main__":
+    run_comparison()

src/dataset.py

@@ -0,0 +1,188 @@
+"""
+Download and prepare TrashNet dataset from GitHub.
+
+Fetches the dataset-resized.zip, converts all images to numpy arrays
+with consistent size, and splits into train/val/test sets.
+
+Saved files:
+    data/processed/X_train.npy, y_train.npy
+    data/processed/X_val.npy,   y_val.npy
+    data/processed/X_test.npy,  y_test.npy
+    data/processed/classes.npy
+"""
+
+import io
+import os
+import sys
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import requests
+import torch
+import yaml
+from PIL import Image
+from sklearn.model_selection import train_test_split
+from torch.utils.data import Dataset
+
+# Add the project root to the python path
+sys.path.append(str(Path(__file__).parent.parent))
+
+
+def load_config(config_path="config.yaml"):
+    with open(config_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+config = load_config()
+CLASSES = config["classes"]
+SPLIT = config["split"]
+IMG_SIZE = tuple(config["img_size"])
+RANDOM_SEED = config["random_seed"]
+
+os.makedirs(name="data", exist_ok=True)
+
+GITHUB_URL = "https://github.com/garythung/trashnet/raw/master/data/dataset-resized.zip"
+RAW_DIR = Path("data/raw")
+OUT_DIR = Path("data/processed")
+
+
+def download_and_extract() -> Path:
+    """
+    Automates data acquisition for reproducibility.
+    Downloads the TrashNet ZIP from GitHub and extracts it to data/raw/.
+    This ensures that anyone running the script gets the exact same starting data.
+    """
+    if (RAW_DIR / "dataset-resized").exists():
+        print("[SKIP] Already extracted.")
+        return RAW_DIR
+
+    print("[DOWNLOAD] TrashNet from GitHub...")
+
+    try:
+        response = requests.get(GITHUB_URL, timeout=120)
+        response.raise_for_status()
+
+        with zipfile.ZipFile(io.BytesIO(response.content)) as zf:
+            zf.extractall(RAW_DIR)
+            print("[OK] Extraction done.")
+    except Exception as e:
+        print(f"[ERROR] Download failed: {e}")
+        raise
+
+    return RAW_DIR
+
+
+def load_images(raw_path: Path) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Data Standardization:
+    Reads all images per class, resizes them to a uniform size (IMG_SIZE),
+    and converts them to RGB. This creates a consistent input format for
+    the neural network, regardless of the original image dimensions or formats.
+    """
+    images, labels = [], []
+
+    for label_index, class_name in enumerate(CLASSES):
+        class_dir = raw_path / "dataset-resized" / class_name
+        if not class_dir.exists():
+            print(f"[WARN] Folder not found: {class_dir}, skipping.")
+            continue
+
+        files = list(class_dir.glob("*.jpg")) + list(class_dir.glob("*.png"))
+        print(f"[LOAD] {class_name}: {len(files)} images")
+
+        for img_path in files:
+            try:
+                img = Image.open(img_path).convert("RGB")
+                img = img.resize(IMG_SIZE)
+                images.append(np.array(img, dtype=np.uint8))
+                labels.append(label_index)
+            except Exception as e:
+                print(f"[WARN] Could not read {img_path}: {e}")
+
+    return np.array(images), np.array(labels, dtype=np.int64)
+
+
+def split_and_save(images: np.ndarray, labels: np.ndarray):
+    """
+    Evaluation Rigor:
+    Splits the data into fixed Train, Validation, and Test sets.
+    - Train: Used to update model weights.
+    - Val: Used to tune hyperparameters and prevent overfitting.
+    - Test: Used for final unbiased evaluation.
+    Saving as .npy files makes loading much faster during training.
+    """
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    train_ratio, val_ratio, test_ratio = SPLIT
+
+    X_train, X_rest, y_train, y_rest = train_test_split(
+        images, labels, test_size=(1 - train_ratio), stratify=labels, random_state=RANDOM_SEED
+    )
+
+    val_size = val_ratio / (val_ratio + test_ratio)
+    X_val, X_test, y_val, y_test = train_test_split(
+        X_rest, y_rest, test_size=(1 - val_size), stratify=y_rest, random_state=RANDOM_SEED
+    )
+
+    splits = {
+        "X_train": X_train,
+        "y_train": y_train,
+        "X_val": X_val,
+        "y_val": y_val,
+        "X_test": X_test,
+        "y_test": y_test,
+    }
+
+    for name, array in splits.items():
+        path = OUT_DIR / f"{name}.npy"
+        np.save(path, array)
+        print(f"[OK] {name}.npy → {array.shape}  dtype={array.dtype}")
+
+    np.save(OUT_DIR / "classes.npy", np.array(CLASSES))
+    print(f"\n[DONE] Splits: " f"Train={len(y_train)} | Val={len(y_val)} | Test={len(y_test)}")
+
+
+class TrashDataset(Dataset):
+    """
+    The Bridge to PyTorch:
+    This class is REQUIRED because PyTorch's DataLoader expects a Dataset object.
+
+    Why this class?
+    1. Efficient Loading: It only loads specific images into RAM when needed (lazy loading).
+    2. Data Augmentation: Allows on-the-fly transformations (rotation, flip, etc.) in __getitem__.
+    3. Tensor Conversion: Handles the conversion from NumPy arrays to PyTorch Tensors.
+    """
+
+    def __init__(self, x_path: Path, y_path: Path, transform=None):
+        """Loads the pre-processed .npy files once into memory."""
+        self.X = np.load(x_path)
+        self.y = torch.from_numpy(np.load(y_path))
+        self.transform = transform
+
+    def __len__(self):
+        """Tells the DataLoader how many samples are in the dataset."""
+        return len(self.X)
+
+    def __getitem__(self, idx):
+        """
+        Fetches a single sample (image + label) at the given index.
+        This is where preprocessing (transforms) happens during training.
+        """
+        img = self.X[idx]
+        label = self.y[idx]
+
+        if self.transform:
+            img = self.transform(img)
+        else:
+            # Default: Convert [H, W, C] (0-255) to [C, H, W] (0.0-1.0) for
+            # PyTorch
+            img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
+
+        return img, label
+
+
+if __name__ == "__main__":
+    raw_path = download_and_extract()
+    images, labels = load_images(raw_path)
+    split_and_save(images, labels)

src/evaluate.py

@@ -0,0 +1,106 @@
+import os
+import sys
+from pathlib import Path
+
+# Add project root to sys.path
+sys.path.append(str(Path(__file__).parent.parent))
+
+import matplotlib  # noqa: E402
+
+matplotlib.use("Agg")  # Use headless backend
+import matplotlib.pyplot as plt  # noqa: E402
+import torch  # noqa: E402
+import torch.nn as nn  # noqa: E402
+import yaml  # noqa: E402
+from sklearn.metrics import (  # noqa: E402
+    ConfusionMatrixDisplay,
+    classification_report,
+    confusion_matrix,
+)
+
+
+def load_config(config_path="config.yaml"):
+    with open(config_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+config = load_config()
+CLASSES = config["classes"]
+
+
+def get_device(config_device):
+    if config_device == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    return config_device
+
+
+DEVICE = get_device(config["device"])
+
+
+def evaluate(model, data_loader, device=DEVICE, save_dir="models/plots"):
+    """
+    Evaluates a PyTorch model on a given DataLoader.
+
+    Args:
+        model: The PyTorch model to evaluate.
+        data_loader: The DataLoader providing the evaluation data.
+        device: The device to run evaluation on (e.g., 'cuda', 'cpu').
+        save_dir: Directory to save plots.
+
+    Returns:
+        avg_loss (float): The average loss over the dataset.
+        accuracy (float): The classification accuracy (0.0 to 1.0).
+    """
+    model.to(device)
+    model.eval()
+    criterion = nn.CrossEntropyLoss()
+
+    total_loss = 0.0
+    correct = 0
+    total = 0
+
+    all_preds = []
+    all_labels = []
+
+    with torch.no_grad():
+        for images, labels in data_loader:
+            images, labels = images.to(device), labels.to(device)
+
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            total_loss += loss.item()
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+            all_preds.extend(predicted.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+
+    avg_loss = total_loss / len(data_loader)
+    accuracy = correct / total
+
+    print("\nEvaluation Results:")
+    print(f"Average Loss: {avg_loss:.4f}")
+    print(f"Accuracy: {accuracy:.4f}")
+
+    # Classification Report
+    print("\nClassification Report:")
+    report = classification_report(
+        all_labels, all_preds, target_names=CLASSES, labels=range(len(CLASSES)), zero_division=0
+    )
+    print(report)
+
+    # Confusion Matrix
+    cm = confusion_matrix(all_labels, all_preds, labels=range(len(CLASSES)))
+    os.makedirs(save_dir, exist_ok=True)
+
+    fig, ax = plt.subplots(figsize=(10, 8))
+    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=CLASSES)
+    disp.plot(cmap=plt.cm.Blues, ax=ax, xticks_rotation=45)
+    plt.title("Confusion Matrix")
+    plt.tight_layout()
+    plt.savefig(f"{save_dir}/confusion_matrix.png")
+    print(f"\nConfusion matrix saved to {save_dir}/confusion_matrix.png")
+
+    return avg_loss, accuracy

src/model.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+from torchvision import models
+from torchvision.models import ResNet18_Weights
+
+
+class SimpleCNN(nn.Module):
+    """
+    A minimalist CNN model as a baseline.
+    Consists of two convolutional layers followed by a fully connected layer.
+    """
+
+    def __init__(self, num_classes=6):
+        super(SimpleCNN, self).__init__()
+        # First Convolutional Block: Takes 3 channels (RGB) as input and
+        # outputs 16
+        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
+        self.relu1 = nn.ReLU()
+        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 224 -> 112
+
+        # Second Convolutional Block
+        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
+        self.relu2 = nn.ReLU()
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  # 112 -> 56
+
+        # Adaptive Pooling ensures the output is always 7x7, regardless of
+        # input size
+        self.adaptive_pool = nn.AdaptiveAvgPool2d((7, 7))
+
+        # Classification Layer
+        self.fc = nn.Linear(32 * 7 * 7, num_classes)
+
+    def forward(self, x):
+        """
+        Defines the forward pass of the data through the network.
+        """
+        x = self.pool1(self.relu1(self.conv1(x)))
+        x = self.pool2(self.relu2(self.conv2(x)))
+        x = self.adaptive_pool(x)
+        x = torch.flatten(x, 1)  # Flatten for the linear layer
+        x = self.fc(x)
+        return x
+
+
+class DeepCNN(nn.Module):
+    """
+    A deeper CNN model with Batch Normalization and Dropout for regularization.
+    Better suited for more complex image features.
+    """
+
+    def __init__(self, num_classes=6):
+        super(DeepCNN, self).__init__()
+
+        # Block 1
+        self.layer1 = nn.Sequential(
+            nn.Conv2d(3, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),  # 112
+        )
+
+        # Block 2
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),  # 56
+        )
+
+        # Block 3
+        self.layer3 = nn.Sequential(
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),  # 28
+        )
+
+        self.adaptive_pool = nn.AdaptiveAvgPool2d((7, 7))
+
+        # Classifier with Dropout to prevent overfitting
+        self.classifier = nn.Sequential(
+            nn.Linear(128 * 7 * 7, 512), nn.ReLU(), nn.Dropout(0.5), nn.Linear(512, num_classes)
+        )
+
+    def forward(self, x):
+        """
+        Forward pass through the sequential layers.
+        """
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.adaptive_pool(x)
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
+
+class ResNet18Transfer(nn.Module):
+    """
+    Transfer Learning model based on ResNet18.
+    Allows loading pretrained weights and freezing the backbone.
+    """
+
+    def __init__(self, num_classes=6, pretrained=True, freeze_backbone=False):
+        super(ResNet18Transfer, self).__init__()
+
+        # Load the ResNet18 model
+        weights = ResNet18_Weights.DEFAULT if pretrained else None
+        self.backbone = models.resnet18(weights=weights)
+
+        # Freeze the backbone if requested
+        if freeze_backbone:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+        # Adjust the final fully connected layer (fc)
+        # ResNet18 fc has 512 input features by default
+        in_features = self.backbone.fc.in_features
+        self.backbone.fc = nn.Linear(in_features, num_classes)
+
+    def forward(self, x):
+        """
+        Uses the ResNet backbone for feature extraction and classification.
+        """
+        return self.backbone(x)

src/train.py

@@ -0,0 +1,240 @@
+import os
+import sys
+from pathlib import Path
+
+# Add project root to sys.path
+sys.path.append(str(Path(__file__).parent.parent))
+
+import matplotlib  # noqa: E402
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+import mlflow  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+import torch.nn as nn  # noqa: E402
+import yaml  # noqa: E402
+from torch.utils.data import DataLoader  # noqa: E402
+from torchvision import transforms  # noqa: E402
+from tqdm import tqdm  # noqa: E402
+
+from src.dataset import TrashDataset  # noqa: E402
+from src.model import DeepCNN, ResNet18Transfer, SimpleCNN  # noqa: E402
+
+
+def load_config(config_path="config.yaml"):
+    with open(config_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+def get_device(config_device):
+    if config_device == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    return config_device
+
+
+class Trainer:
+    """
+    Handles the training and validation lifecycle of a model with MLflow tracking.
+    """
+
+    def __init__(self, model, train_loader, val_loader, config, model_name):
+        self.config = config
+        self.model_name = model_name
+        self.device = get_device(config["device"])
+        self.model = model.to(self.device)
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.epochs = config["epochs"]
+        self.patience = config.get("patience", 5)
+
+        # Handle class imbalance with weights
+        y_train = np.load("data/processed/y_train.npy")
+        class_counts = np.bincount(y_train)
+        weights = 1.0 / class_counts
+        weights = torch.FloatTensor(weights).to(self.device)
+
+        self.criterion = nn.CrossEntropyLoss(weight=weights)
+        self.optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"])
+        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            self.optimizer, T_max=self.epochs
+        )
+        self.history = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}
+        self.checkpoint_path = f"models/{model_name.lower()}_best.pth"
+
+    def train_epoch(self):
+        self.model.train()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+
+        pbar = tqdm(self.train_loader, desc="Training", leave=False)
+        for images, labels in pbar:
+            images, labels = images.to(self.device), labels.to(self.device)
+
+            self.optimizer.zero_grad()
+            outputs = self.model(images)
+            loss = self.criterion(outputs, labels)
+            loss.backward()
+            self.optimizer.step()
+
+            running_loss += loss.item()
+            _, predicted = torch.max(outputs, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+            pbar.set_postfix({"loss": f"{loss.item():.4f}", "acc": f"{correct/total:.4f}"})
+
+        return running_loss / len(self.train_loader), correct / total
+
+    def validate(self):
+        self.model.eval()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+
+        with torch.no_grad():
+            for images, labels in self.val_loader:
+                images, labels = images.to(self.device), labels.to(self.device)
+                outputs = self.model(images)
+                loss = self.criterion(outputs, labels)
+
+                running_loss += loss.item()
+                _, predicted = torch.max(outputs, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+
+        return running_loss / len(self.val_loader), correct / total
+
+    def plot_history(self):
+        save_path = f"models/plots/{self.model_name.lower()}_history.png"
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        epochs_range = range(1, len(self.history["train_loss"]) + 1)
+
+        plt.figure(figsize=(12, 5))
+
+        # Plot Loss
+        plt.subplot(1, 2, 1)
+        plt.plot(epochs_range, self.history["train_loss"], label="Train Loss")
+        plt.plot(epochs_range, self.history["val_loss"], label="Val Loss")
+        plt.title(f"{self.model_name} - Loss")
+        plt.xlabel("Epochs")
+        plt.ylabel("Loss")
+        plt.legend()
+
+        # Plot Accuracy
+        plt.subplot(1, 2, 2)
+        plt.plot(epochs_range, self.history["train_acc"], label="Train Acc")
+        plt.plot(epochs_range, self.history["val_acc"], label="Val Acc")
+        plt.title(f"{self.model_name} - Accuracy")
+        plt.xlabel("Epochs")
+        plt.ylabel("Accuracy")
+        plt.legend()
+
+        plt.tight_layout()
+        plt.savefig(save_path)
+        print(f"--> Training history plot saved to {save_path}")
+        mlflow.log_artifact(save_path)
+
+    def fit(self):
+        mlflow.set_experiment("Trash Classifier")
+        with mlflow.start_run(run_name=self.model_name):
+            mlflow.log_params(self.config)
+            mlflow.log_param("model_architecture", self.model_name)
+
+            print(f"\nStarting training for {self.model_name} on {self.device}...")
+            best_val_acc = 0.0
+            epochs_no_improve = 0
+
+            for epoch in range(self.epochs):
+                train_loss, train_acc = self.train_epoch()
+                val_loss, val_acc = self.validate()
+                self.scheduler.step()
+
+                self.history["train_loss"].append(train_loss)
+                self.history["train_acc"].append(train_acc)
+                self.history["val_loss"].append(val_loss)
+                self.history["val_acc"].append(val_acc)
+
+                mlflow.log_metric("train_loss", train_loss, step=epoch)
+                mlflow.log_metric("train_acc", train_acc, step=epoch)
+                mlflow.log_metric("val_loss", val_loss, step=epoch)
+                mlflow.log_metric("val_acc", val_acc, step=epoch)
+                mlflow.log_metric("lr", self.optimizer.param_groups[0]["lr"], step=epoch)
+
+                print(
+                    f"Epoch [{epoch + 1}/{self.epochs}] "
+                    f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | "
+                    f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f} | "
+                    f"LR: {self.optimizer.param_groups[0]['lr']:.6f}"
+                )
+
+                if val_acc > best_val_acc:
+                    best_val_acc = val_acc
+                    epochs_no_improve = 0
+                    os.makedirs("models", exist_ok=True)
+                    torch.save(self.model.state_dict(), self.checkpoint_path)
+                    print(f"--> Saved best model for {self.model_name} with Val Acc: {val_acc:.4f}")
+                    mlflow.log_artifact(self.checkpoint_path)
+                else:
+                    epochs_no_improve += 1
+                    if epochs_no_improve >= self.patience:
+                        print(f"Early stopping triggered after {epoch + 1} epochs.")
+                        break
+
+            self.plot_history()
+        return self.history
+
+
+if __name__ == "__main__":
+    config = load_config()
+
+    train_transform = transforms.Compose(
+        [
+            transforms.ToPILImage(),
+            transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
+            transforms.RandomHorizontalFlip(),
+            transforms.RandomRotation(15),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    val_transform = transforms.Compose(
+        [
+            transforms.ToPILImage(),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    processed_dir = Path("data/processed")
+    if not (processed_dir / "X_train.npy").exists():
+        print("Error: Processed data not found. Please run src/dataset.py first.")
+    else:
+        train_ds = TrashDataset(
+            processed_dir / "X_train.npy", processed_dir / "y_train.npy", transform=train_transform
+        )
+        val_ds = TrashDataset(
+            processed_dir / "X_val.npy", processed_dir / "y_val.npy", transform=val_transform
+        )
+
+        train_loader = DataLoader(train_ds, batch_size=config["batch_size"], shuffle=True)
+        val_loader = DataLoader(val_ds, batch_size=config["batch_size"], shuffle=False)
+
+        num_classes = len(config["classes"])
+        models_to_train = {
+            "SimpleCNN": SimpleCNN(num_classes=num_classes),
+            "DeepCNN": DeepCNN(num_classes=num_classes),
+            "ResNet18": ResNet18Transfer(num_classes=num_classes, pretrained=True),
+        }
+
+        for name, model in models_to_train.items():
+            trainer = Trainer(model, train_loader, val_loader, config, name)
+            trainer.fit()
+
+        print("\nAll models trained. Starting comparison...")
+        from src.comparison import run_comparison
+
+        run_comparison()