added more py files for build import

data.py

@@ -0,0 +1,84 @@
+# -------------------------------------------------------------
+# This script handles ALL dataset-related logic:
+# 1. Builds image transformations for training and testing
+# 2. Loads datasets from "train", "val", "test" folders
+# 3. Creates DataLoaders for batching
+# 4. Exports label names to "models/labels.json" for later use
+# -------------------------------------------------------------
+
+from pathlib import Path
+import os, json, torch
+from typing import Dict, Tuple
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
+# ------------------------------
+# BUILD TRANSFORMS (AUGMENTATION)
+# ------------------------------
+def build_transforms(img_size: int = 224) -> Tuple[transforms.Compose, transforms.Compose]:
+
+    # --- Training transforms ---
+    # These help the model generalize better by creating variations of input images
+    train_tf = transforms.Compose([
+        transforms.Resize((img_size, img_size)),
+        transforms.RandomHorizontalFlip(),
+        transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
+        transforms.RandomRotation(45, fill=(0,)),  # Background filled with black
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406],
+                             [0.229, 0.224, 0.225]),
+    ])
+
+    # --- Evaluation transforms ---
+    # Only resize + normalize (no randomness to keep results consistent)
+    eval_tf = transforms.Compose([
+        transforms.Resize((img_size, img_size)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406],
+                             [0.229, 0.224, 0.225]),
+    ])
+    return train_tf, eval_tf
+
+
+# ------------------------------
+# LOAD DATASETS & BUILD LOADERS
+# ------------------------------
+# Load datasets and create data loaders
+def load_data(data_root: str, bs: int, num_workers: int = 2) -> Dict:
+
+    # Get transformation pipelines
+    train_tf, eval_tf = build_transforms(224)
+
+    # --- Load image folders ---
+    # Each folder should look like:
+    # data_root/train/class_name/*.jpg
+    # data_root/val/class_name/*.jpg
+    # data_root/test/class_name/*.jpg
+    dsets = {
+        "train": datasets.ImageFolder(os.path.join(data_root, "train"), transform=train_tf),
+        "val":   datasets.ImageFolder(os.path.join(data_root, "val"),   transform=eval_tf),
+        "test":  datasets.ImageFolder(os.path.join(data_root, "test"),  transform=eval_tf),
+    }
+
+    # --- Save label mapping to file ---
+    # This creates a file: models/labels.json
+    # Example content: {"0": "paper", "1": "plastic", "2": "metal"}
+    Path("models").mkdir(parents=True, exist_ok=True)
+    with open("models/labels.json", "w") as f:
+        json.dump({v: k for k, v in dsets["train"].class_to_idx.items()}, f, indent=2)
+
+    # --- Check if GPU is available ---
+    # If yes, use "pinned memory" to load faster
+    pin = torch.cuda.is_available()
+
+    # --- Build DataLoaders ---
+    # DataLoader helps in batching, shuffling, and parallel loading
+    loaders = {
+        "train": DataLoader(dsets["train"], batch_size=bs, shuffle=True,  num_workers=num_workers, pin_memory=pin),
+        "val":   DataLoader(dsets["val"],   batch_size=bs, shuffle=False, num_workers=num_workers, pin_memory=pin),
+        "test":  DataLoader(dsets["test"],  batch_size=bs, shuffle=False, num_workers=num_workers, pin_memory=pin),
+    }
+
+    # Return both datasets and loaders for training scripts
+    return {"dsets": dsets, "loaders": loaders}

dataset.py

@@ -0,0 +1,111 @@
+# -------------------------------------------------------------
+# This script:
+# Downloads the Trash-Type Image Dataset from Kaggle
+# Keeps only the classes: paper, plastic, metal (groups others)
+# Splits data into train / val / test
+# Copies them into the folder structure:
+#   data/images/train/paper/
+#   data/images/val/plastic/
+#   data/images/test/metal/
+# -------------------------------------------------------------
+
+import os, shutil, random
+from pathlib import Path
+import kagglehub
+
+# ---- CLASS FILTERING ----
+CLASS_KEEP = {"paper", "plastic", "metal"}      # Classes we care about
+TO_OTHERS  = {"glass", "cardboard", "trash"}    # Classes to group as "others"
+
+# ---- TARGET DIRECTORY ----
+TARGET = Path("data/images") # Where to save the new dataset
+
+# ---- DATA SPLIT RATIOS ----
+SPLITS = {"train":0.7,"val":0.15,"test":0.15}
+
+# ---- IMAGE EXTENSIONS TO LOOK FOR ----
+IMG_EXTS = (".jpg",".jpeg",".png",".bmp",".webp")
+
+# ---- SET FIXED RANDOM SEED (for reproducibility) ----
+random.seed(42)
+
+# -------------------------------------------------------------
+# STEP 1: INFER CLASS FROM FILE PATH
+# -------------------------------------------------------------
+def infer_class(path: Path) -> str:
+    parts = [p.lower() for p in path.parts]
+    for p in reversed(parts): # Search backwards in path parts
+        if p in CLASS_KEEP:
+            return p
+        if p in TO_OTHERS:
+            return "others"
+    return "others"  # Fallback if no match found
+
+# -------------------------------------------------------------
+# STEP 2: FIND ALL IMAGES IN THE DOWNLOADED DATASET
+# -------------------------------------------------------------
+def gather_images(root: Path):
+    imgs = []
+    for fp in root.rglob("*"): # Recursively walk through subfolders
+        if fp.suffix.lower() in IMG_EXTS:
+            cls = infer_class(fp.parent)
+            imgs.append((fp, cls))
+    return imgs
+
+# -------------------------------------------------------------
+# STEP 3: SPLIT INTO TRAIN/VAL/TEST AND COPY FILES
+# -------------------------------------------------------------
+def split_and_copy(items):
+
+    by_cls = {}
+
+    # Group all images by class
+    for fp, cls in items:
+        by_cls.setdefault(cls, []).append(fp)
+
+    TARGET.mkdir(parents=True, exist_ok=True)
+
+    # For each class, randomly split its images
+    for cls, files in by_cls.items():
+
+        random.shuffle(files)
+        n = len(files); n_tr = int(SPLITS["train"]*n); n_va = int(SPLITS["val"]*n)
+
+        # Split into train, val, test
+        splits = {"train":files[:n_tr], "val":files[n_tr:n_tr+n_va], "test":files[n_tr+n_va:]}
+
+        # Copy images to correct folders
+        for split, fps in splits.items():
+            outdir = TARGET / split / cls
+            outdir.mkdir(parents=True, exist_ok=True)
+            for src in fps:
+                shutil.copy(src, outdir / src.name)
+
+        # Print summary
+        print(f"{cls}: {n} -> train {len(splits['train'])}, val {len(splits['val'])}, test {len(splits['test'])}")
+
+# -------------------------------------------------------------
+# STEP 4: MAIN FUNCTION – DOWNLOAD & PREPARE DATASET
+# -------------------------------------------------------------
+def main():
+
+    # Download dataset (cached locally after first time)
+    root = Path(kagglehub.dataset_download("farzadnekouei/trash-type-image-dataset"))
+    print("Dataset at:", root)
+
+    # Collect all image paths + inferred labels
+    items = gather_images(root)
+
+    if not items:
+        raise SystemExit("No images found. Check dataset layout.")
+    print("Classes detected:", sorted({c for _, c in items}))
+
+    # Split and copy into new structure
+    split_and_copy(items)
+    print("Flattened dataset ready in:", TARGET)
+
+# -------------------------------------------------------------
+# ENTRY POINT (only runs when executed directly)
+# -------------------------------------------------------------
+if __name__ == "__main__":
+    main()

eval.py

@@ -0,0 +1,108 @@
+# -------------------------------------------------------------
+# This script evaluates your trained model:
+# Runs inference on the test dataset
+# Calculates accuracy, precision, recall, F1-score
+# Plots a confusion matrix
+# Saves everything in the /reports folder
+# -------------------------------------------------------------
+
+import json
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+import torch
+from sklearn.metrics import classification_report, confusion_matrix
+import matplotlib.pyplot as plt
+
+# -------------------------------------------------------------
+# HELPER FUNCTION: Plot Confusion Matrix
+# -------------------------------------------------------------
+def _plot_confusion_matrix(cm: np.ndarray, class_names: List[str], title: str, outpath: str = "reports/confusion_matrix.png"):
+
+    cm = np.array(cm)  # Ensure it’s a NumPy array
+
+     # Create figure and axes for plotting
+    fig, ax = plt.subplots(figsize=(7, 6))
+    im = ax.imshow(cm) # Display matrix as an image (heatmap)
+
+    # Label the x and y axes with class names
+    ax.set_xticks(range(len(class_names)))
+    ax.set_yticks(range(len(class_names)))
+    ax.set_xticklabels(class_names, rotation=45, ha="right")
+    ax.set_yticklabels(class_names)
+    ax.set_xlabel("Predicted label")
+    ax.set_ylabel("True label")
+    ax.set_title(title)
+
+    # Add numeric values inside each cell
+    for i in range(cm.shape[0]):
+        for j in range(cm.shape[1]):
+            ax.text(j, i, str(cm[i, j]), ha="center", va="center")
+
+    # Add colorbar to show intensity scale
+    cbar = fig.colorbar(im, ax=ax)
+    cbar.ax.set_ylabel("counts") # Label for the colorbar
+
+    fig.tight_layout()
+
+    # Save the figure to file
+    Path(outpath).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(outpath, dpi=200)
+    plt.close(fig)
+
+# -------------------------------------------------------------
+# MAIN FUNCTION: Evaluate Model on Test Data
+# -------------------------------------------------------------
+def eval_on_test(model, loader, class_names: List[str], device) -> Dict:
+
+    model.eval() # Switch model to evaluation mode (disables dropout/batchnorm)
+    y_true, y_pred = [], [] # Store true and predicted labels
+
+    # Disable gradient computation (faster inference)
+    with torch.no_grad():
+        for xb, yb in loader:                       # Loop through test batches     
+            xb = xb.to(device, non_blocking=True)   # Move data to device
+            out = model(xb).cpu()                   # Forward pass and move output to CPU 
+            y_true += yb.tolist()                   # Collect true labels
+            y_pred += out.argmax(1).tolist()        # Collect predicted labels
+
+    # ---- PRINT TEXT REPORT ----
+    print("\n=== Classification Report ===")
+    print(classification_report(y_true, y_pred, target_names=class_names, digits=4))
+
+    # ---- CONFUSION MATRIX ----
+    labels = list(range(len(class_names))) # Numeric labels [0, 1, 2, ...]
+    cm = confusion_matrix(y_true, y_pred, labels=labels)
+    print("Confusion Matrix:\n", cm)
+    
+    # ---- CALCULATE OVERALL ACCURACY ----
+    total = cm.sum()
+    acc = float(cm.trace()) / float(total) if total > 0 else 0.0 # Correct predictions / Total predictions
+
+    # ---- SAVE METRICS INTO A DICTIONARY ----
+    metrics = {
+        "accuracy": acc,
+        "labels": class_names,
+        "classification_report": classification_report(
+            y_true, y_pred, target_names=class_names, digits=4, output_dict=True
+        ),
+        "confusion_matrix": cm.tolist(),
+    }
+
+    # ---- SAVE METRICS TO FILE ----
+    Path("reports").mkdir(parents=True, exist_ok=True)
+    with open("reports/metrics.json", "w") as f:
+        json.dump(metrics, f, indent=2)
+
+    # ---- PLOT AND SAVE CONFUSION MATRIX IMAGE ----
+    _plot_confusion_matrix(
+        cm,
+        class_names,
+        title="Recycle Classifier Confusion Matrix",
+        outpath="reports/confusion_matrix.png",
+    )
+
+    print("Saved: reports/metrics.json and reports/confusion_matrix.png")
+    return metrics
+

explain.py

@@ -0,0 +1,110 @@
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from data import build_transforms
+from model import build_model
+import json
+from pathlib import Path
+
+
+def generate_gradcam(img, model, device, class_names, threshold=0.5):
+
+    """
+    Generate Grad-CAM for a given image and model.
+
+    Args:
+        img (PIL.Image): Input image.
+        model (torch.nn.Module): Pretrained model.
+        device (torch.device): Torch device (cpu/cuda).
+        class_names (list): List of class labels.
+        threshold (float): Threshold for masking heatmap.
+
+    Returns:
+        overlay_rgb (PIL.Image): Overlay of heatmap on original image.
+        heatmap (np.ndarray): Heatmap array.
+        pred_label (str): Predicted class label.
+        confidence (float): Probability of predicted class.
+    """
+
+    # Load transforms
+    _, eval_transforms = build_transforms(224)
+    
+    # Handle input type
+    if isinstance(img, str):
+        img = Image.open(img).convert("RGB")
+    elif not isinstance(img, Image.Image):
+        raise ValueError("Input must be a file path or PIL.Image")
+    
+    img_tensor = eval_transforms(img).unsqueeze(0).to(device)
+
+    # Hook storage
+    feats, grads = [], []
+    def fwd_hook(m, i, o): feats.append(o)
+    def bwd_hook(m, gi, go): grads.append(go[0])
+
+    # Register hooks
+    layer = model.layer4[1].conv2
+    layer.register_forward_hook(fwd_hook)
+    layer.register_full_backward_hook(bwd_hook)
+
+    # Forward pass
+    out = model(img_tensor)
+    pred_idx = out.argmax(1).item()
+    score = out[0, pred_idx]
+
+    # Confidence (softmax)
+    probs = torch.softmax(out, dim=1)
+    confidence = probs[0, pred_idx].item()
+
+    # Backward pass
+    model.zero_grad()
+    score.backward()
+
+    # Grad-CAM calculation
+    grad = grads[0][0].detach().cpu().numpy()
+    feat = feats[0][0].detach().cpu().numpy()
+    weights = grad.mean(axis=(1, 2))
+    cam = np.maximum(np.sum(weights[:, None, None] * feat, axis=0), 0)
+    cam = cv2.resize(cam, img.size)
+    cam = cam / cam.max()
+
+    # Mask
+    mask = (cam > threshold).astype(np.uint8) * 255
+
+    # Heatmap
+    heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_TURBO)
+    heatmap_masked = cv2.bitwise_and(heatmap, heatmap, mask=mask)
+    heatmap_rgb = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+
+    # Overlay
+    orig_bgr = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+    overlay = cv2.addWeighted(orig_bgr, 0.5, heatmap_masked, 0.5, 0)
+    overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
+
+    return Image.fromarray(overlay_rgb), heatmap_rgb, class_names[pred_idx], confidence
+
+
+if __name__ == "__main__":
+    
+    # Load labels
+    with open("models/labels.json", "r") as f:
+        idx2name = {int(k): v for k, v in json.load(f).items()}
+    class_names = [idx2name[i] for i in sorted(idx2name.keys())]
+
+    # Load model
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = build_model(num_classes=len(class_names), freeze_backbone=False, device=device)
+    model.load_state_dict(torch.load("models/resnet18_best.pt", map_location=device))
+    model.eval()
+    print("Model loaded.")
+
+    # Test with one image
+    img = Image.open("data/own_images/test/paper/paper_normal_dark_23sec_013.jpg").convert("RGB")
+    overlay, heatmap, pred_label, conf = generate_gradcam(img, model, device, class_names)
+
+    print(f"Final Prediction: {pred_label} ({conf:.2%})")
+    overlay.show()  # Display overlay
+    heatmap_img = Image.fromarray(heatmap)
+    heatmap_img.show()  # Display heatmap
+    

images_combined.py

@@ -0,0 +1,105 @@
+# -------------------------------------------------------------
+# This script merges the two datasets:
+#   1️⃣ Kaggle dataset (data/images)
+#   2️⃣ Your own collected images (data/own_images)
+#
+# It combines both into one new dataset folder:
+# data/combined/
+#
+# The structure after merging looks like:
+#   data/combined/
+#       ├── train/
+#       │    ├── paper/
+#       │    ├── plastic/
+#       │    └── metal/
+#       ├── val/
+#       └── test/
+#
+# So the model can train on all available images together.
+# -------------------------------------------------------------
+
+from pathlib import Path
+import shutil
+
+# ---- Define source and destination paths ----
+SRC_K = Path("data/images")        # Kaggle-prepared dataset
+SRC_O = Path("data/own_images")    # Your own collected images
+DST   = Path("data/combined")      # Combined output dataset
+
+# -------------------------------------------------------------
+# FUNCTION: copy_all
+# -------------------------------------------------------------
+def copy_all(src_dir: Path, dst_dir: Path):
+    
+    if not src_dir.exists():
+        return 0 # Skip if source folder doesn’t exist
+    
+    dst_dir.mkdir(parents=True, exist_ok=True) # Make sure target exists
+    n = 0
+
+    # Copy all files from source to destination
+    for p in src_dir.glob("*"):
+        if p.is_file(): # Ignore subfolders
+            shutil.copy(p, dst_dir / p.name)
+            n += 1
+    return n # Return count of copied files
+
+# -------------------------------------------------------------
+# FUNCTION: classes_in
+# -------------------------------------------------------------
+def classes_in(root: Path):
+    out = set()
+    for split in ["train", "val", "test"]:
+        d = root / split
+        if d.exists():
+            for c in d.iterdir():
+                if c.is_dir():
+                    out.add(c.name)
+    return out
+
+# -------------------------------------------------------------
+# STEP 1: Identify all classes across Kaggle + own datasets
+# -------------------------------------------------------------
+classes = sorted(classes_in(SRC_K) | classes_in(SRC_O)) # Combine both sets
+
+if not classes:
+    raise SystemExit("No classes found in data/images or data/own_images. Build them first.")
+
+print("Classes to combine:", classes)
+
+# -------------------------------------------------------------
+# STEP 2: Create folder structure for combined dataset
+# -------------------------------------------------------------
+for split in ["train", "val", "test"]:
+    for cls in classes:
+        (DST / split / cls).mkdir(parents=True, exist_ok=True)
+
+# -------------------------------------------------------------
+# STEP 3: Combine TRAIN images (both Kaggle + Own)
+# -------------------------------------------------------------
+for cls in classes:
+    n1 = copy_all(SRC_K / "train" / cls, DST / "train" / cls) # Copy from Kaggle
+    n2 = copy_all(SRC_O / "train" / cls, DST / "train" / cls) # Copy from own
+    print(f"train/{cls}: +{n1} (kaggle) +{n2} (own)")
+
+# -------------------------------------------------------------
+# STEP 4: Combine VALIDATION and TEST images
+# -------------------------------------------------------------
+for split in ["val", "test"]:
+
+    for cls in classes:
+
+        n = copy_all(SRC_K / split / cls, DST / split / cls) # Copy from Kaggle first
+
+        # If none, fill from own images 
+        if n == 0:  
+            n = copy_all(SRC_O / split / cls, DST / split / cls)
+            if n > 0:
+                print(f"{split}/{cls}: filled from own ({n} files)")
+        else:
+            print(f"{split}/{cls}: +{n} (kaggle)")
+
+# -------------------------------------------------------------
+# STEP 5: Done — print where the merged dataset is saved
+# -------------------------------------------------------------
+print("Combined dataset at:", DST)

main.py

@@ -0,0 +1,99 @@
+# -------------------------------------------------------------
+# HOW TO RUN THIS SCRIPT:
+# -------------------------------------------------------------
+# 1️⃣ Create a virtual environment:
+#     python -m venv .venv && .\.venv\Scripts\Activate
+#
+# 2️⃣ Install dependencies:
+#     pip install -r requirements.txt
+#
+# 3️⃣ Train and evaluate model:
+#     python src/main.py --data data/combined --epochs 15 --lr 3e-4 --bs 32
+# -------------------------------------------------------------
+
+import argparse, random
+from pathlib import Path
+from typing import Dict
+import numpy as np
+import torch
+
+from data import load_data          
+from model import build_model       
+from train import train_model       
+from eval import eval_on_test       
+
+# -------------------------------------------------------------
+# FUNCTION: set_seed
+# -------------------------------------------------------------
+def set_seed(seed: int = 42):
+    random.seed(seed); np.random.seed(seed)
+    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+# -------------------------------------------------------------
+# FUNCTION: make_dirs
+# -------------------------------------------------------------
+def make_dirs():
+    Path("models").mkdir(parents=True, exist_ok=True)
+    Path("reports").mkdir(parents=True, exist_ok=True)
+
+# -------------------------------------------------------------
+# FUNCTION: get_device
+# -------------------------------------------------------------
+def get_device():
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# -------------------------------------------------------------
+# FUNCTION: parse_args
+# -------------------------------------------------------------
+def parse_args():
+    p = argparse.ArgumentParser(description="Recycle Material Classifier (ResNet-18)")
+    p.add_argument("--data", type=str, default="data/images", help="root with train/val/test")
+    p.add_argument("--epochs", type=int, default=10)
+    p.add_argument("--bs", type=int, default=32)
+    p.add_argument("--lr", type=float, default=1e-3)
+    p.add_argument("--freeze", action="store_true", help="freeze ResNet backbone")
+    p.add_argument("--seed", type=int, default=42)
+    return p.parse_args()
+
+# -------------------------------------------------------------
+# FUNCTION: main
+# -------------------------------------------------------------
+def main():
+
+    # 1️⃣ Read input arguments
+    args = parse_args()
+
+    # 2️⃣ Ensure results are reproducible
+    set_seed(args.seed)
+
+    # 3️⃣ Create folders for saving models/reports
+    make_dirs()
+
+    # 4️⃣ Detect whether to use GPU or CPU
+    device = get_device()
+    print("Device:", device)
+
+    # 5️⃣ Load dataset and dataloaders
+    out = load_data(args.data, args.bs)
+    dsets, loaders = out["dsets"], out["loaders"]
+    # Extract class names in correct order
+    class_names = [k for k, _ in sorted(dsets["train"].class_to_idx.items(), key=lambda x: x[1])]
+    print("Classes:", class_names)
+
+    # 6️⃣ Build model (ResNet18 with custom classification head)
+    model = build_model(num_classes=len(class_names), freeze_backbone=args.freeze, device=device)
+
+    # 7️⃣ Train model using training and validation sets
+    model, best_val = train_model(model, loaders, epochs=args.epochs, lr=args.lr, device=device)
+    print(f"Best val acc: {best_val:.4f}")
+
+    # 8️⃣ Evaluate final model on the test set and save reports
+    eval_on_test(model, loaders["test"], class_names, device)
+
+# -------------------------------------------------------------
+# ENTRY POINT (runs when executed directly)
+# -------------------------------------------------------------
+if __name__ == "__main__":
+    main()

own_images.py

@@ -0,0 +1,65 @@
+# -------------------------------------------------------------
+# This script prepares your OWN collected images
+# (separate from the Kaggle dataset) by splitting them into:
+# train / val / test folders
+#
+# Example output structure after running this:
+# data/own_images/
+# ├── train/
+# │    ├── paper/
+# │    ├── plastic/
+# │    └── metal/
+# ├── val/
+# └── test/
+# -------------------------------------------------------------
+
+from pathlib import Path
+import shutil, random
+
+# ---- SOURCE & DESTINATION FOLDERS ----
+SRC = Path("data/own_images")  # Folder where your raw images are currently stored
+DST = Path("data/own_images")  # Output folder (same location, will create train/val/test inside)
+
+# ---- SPLIT RATIOS ----
+SPLIT = {"train":0.70, "val":0.15, "test":0.15}
+
+# ---- SUPPORTED IMAGE EXTENSIONS ----
+EXTS = {".jpg",".jpeg",".png",".bmp",".webp"}
+
+# ---- FIX RANDOMNESS (for reproducible splits) ----
+random.seed(42)
+
+# -------------------------------------------------------------
+# FUNCTION: gather()
+# -------------------------------------------------------------
+def gather(p): return [f for f in p.rglob("*") if f.suffix.lower() in EXTS]
+
+# -------------------------------------------------------------
+# FUNCTION: split_copy()
+# -------------------------------------------------------------
+def split_copy(cls):
+
+    # 1️⃣ Gather all image paths for this class
+    files = gather(SRC/cls); random.shuffle(files) # Shuffle before splitting for randomness
+   
+    # 2️⃣ Compute how many go into each split
+    n=len(files); n_tr=int(n*SPLIT["train"]); n_va=int(n*SPLIT["val"])
+    
+    # Slice list into 3 parts
+    parts = {"train":files[:n_tr], "val":files[n_tr:n_tr+n_va], "test":files[n_tr+n_va:]}
+
+    # 3️⃣ Copy files into their corresponding split folders
+    for split, arr in parts.items():
+        out = DST/split/cls; out.mkdir(parents=True, exist_ok=True) # Create directory if missing
+        for src in arr: shutil.copy(src, out/src.name) # Copy image file
+
+    # 4️⃣ Print summary for this class
+    print(f"{cls}: {n} -> {len(parts['train'])}/{len(parts['val'])}/{len(parts['test'])}")
+
+# -------------------------------------------------------------
+# MAIN EXECUTION: process all 3 recyclable material classes
+# -------------------------------------------------------------
+for cls in ["paper","plastic","metal"]:
+    split_copy(cls)
+
+print("Done", DST)

tempCodeRunnerFile.py

@@ -0,0 +1,4 @@
+        webcam_input = gr.Camera(label="Take a picture", type="pil")
+            json_output2 = gr.JSON(label="Prediction (top class + confidence %)") 
+            label_output2 = gr.Label(num_top_classes=3, label="Top-3 probabilities")
+            webcam_input.change(predict, inputs=webcam_input, outputs=[json_output2, label_output2])

train.py

@@ -0,0 +1,244 @@
+# -------------------------------------------------------------
+# This script handles the full training loop:
+# 1️⃣ Train and validate a model for multiple epochs
+# 2️⃣ Track accuracy, loss, learning rate, etc.
+# 3️⃣ Save the best model (highest val accuracy)
+# 4️⃣ Automatically stop early if val accuracy stops improving
+# 5️⃣ Save training history and graphs
+# -------------------------------------------------------------
+
+import json, time
+from pathlib import Path
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from typing import Tuple, Dict, List
+from torch.utils.tensorboard import SummaryWriter
+from tqdm.auto import tqdm
+
+# -------------------------------------------------------------
+# HELPER: Compute batch accuracy
+# -------------------------------------------------------------
+def accuracy_from_logits(logits: torch.Tensor, y: torch.Tensor) -> float:
+    preds = torch.argmax(logits, dim=1)
+    return (preds == y).float().mean().item()
+
+# -------------------------------------------------------------
+# HELPER: Format seconds to mm:ss (for clean epoch timing)
+# -------------------------------------------------------------
+def _fmt_time(s: float) -> str:
+    m, s = divmod(int(s), 60)
+    return f"{m:02d}:{s:02d}"
+
+# -------------------------------------------------------------
+# HELPER: Get all target labels from a dataset
+# -------------------------------------------------------------
+def _get_targets(ds) -> torch.Tensor:
+    if hasattr(ds, "targets"):
+        return torch.as_tensor(ds.targets)
+    return torch.as_tensor([t for _, t in ds.samples])
+
+# -------------------------------------------------------------
+# HELPER: Count how many samples per class
+# -------------------------------------------------------------
+def _class_counts(ds) -> List[str]:
+    if not hasattr(ds, "classes"): return []
+    y = _get_targets(ds)
+    counts = torch.bincount(y, minlength=len(ds.classes)).tolist()
+    return [f"{cls}({cnt})" for cls, cnt in zip(ds.classes, counts)]
+
+# -------------------------------------------------------------
+# HELPER: Count total and trainable parameters
+# -------------------------------------------------------------
+def _num_params(model) -> Dict[str, int]:
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return {"total": total, "trainable": trainable}
+
+# -------------------------------------------------------------
+# HELPER: Show current learning rate(s)
+# -------------------------------------------------------------
+def _current_lrs(optimizer) -> str:
+    lrs = sorted({pg["lr"] for pg in optimizer.param_groups})
+    if len(lrs) == 1: return f"{lrs[0]:.2e}"
+    return ", ".join(f"{lr:.2e}" for lr in lrs)
+
+# -------------------------------------------------------------
+# FUNCTION: Run one full epoch (Train OR Validation)
+# -------------------------------------------------------------
+def run_epoch(model,loader,criterion,optimizer,device,train: bool,epoch: int,epochs: int,) -> Tuple[float, float]:
+    
+    model.train() if train else model.eval()
+    epoch_loss = 0.0
+    epoch_acc = 0.0
+    total = 0
+
+    # tqdm creates a nice progress bar in terminal
+    bar = tqdm(
+        loader,
+        desc=f"[{epoch:02d}/{epochs}] {'train' if train else 'val  '}",
+        leave=False,
+        dynamic_ncols=True,
+    )
+
+    for xb, yb in bar:
+        xb = xb.to(device, non_blocking=True)
+        yb = yb.to(device, non_blocking=True)
+
+        if train:
+            optimizer.zero_grad()
+
+        with torch.set_grad_enabled(train):
+
+            out = model(xb)             # Forward pass 
+            loss = criterion(out, yb)   # Compute loss
+
+            if train:
+                loss.backward()         # Backpropagation 
+                optimizer.step()        # Update weights
+
+        # Update running totals
+        bsz = xb.size(0)
+        epoch_loss += loss.item() * bsz
+        epoch_acc  += accuracy_from_logits(out, yb) * bsz
+        total      += bsz
+
+        # Update progress bar text
+        bar.set_postfix(loss=f"{epoch_loss/max(total,1):.4f}",
+                        acc=f"{epoch_acc/max(total,1):.4f}")
+        
+    # Return average loss and accuracy
+    return epoch_loss / total, epoch_acc / total
+
+
+# -------------------------------------------------------------
+# FUNCTION: Save training curves (loss + accuracy)
+# -------------------------------------------------------------
+def _save_final_figure(history, best_val_acc, outpath="reports/curves_final.png", title="ResNet (Fine-tune)"):
+
+    Path(outpath).parent.mkdir(parents=True, exist_ok=True)
+    xs = range(1, len(history["train_loss"]) + 1)
+
+    fig, axs = plt.subplots(1, 2, figsize=(11, 4.5))
+
+    # ---- Loss subplot ----
+    axs[0].plot(xs, history["train_loss"], label="train")
+    axs[0].plot(xs, history["val_loss"],   label="val")
+    axs[0].set_title("Loss"); axs[0].set_xlabel("Epoch"); axs[0].set_ylabel("Loss"); axs[0].legend()
+
+     # ---- Accuracy subplot ----
+    axs[1].plot(xs, history["train_acc"], label="train")
+    axs[1].plot(xs, history["val_acc"],   label="val")
+    axs[1].set_title("Accuracy"); axs[1].set_xlabel("Epoch"); axs[1].set_ylabel("Acc"); axs[1].legend()
+
+    fig.suptitle(f"{title} | best val acc = {best_val_acc:.4f}")
+    fig.tight_layout(rect=[0, 0, 1, 0.95])
+    plt.savefig(outpath, dpi=200)
+    plt.close(fig)
+
+    # Save training history as JSON
+    with open("reports/history.json", "w") as f:
+        json.dump(history, f, indent=2)
+
+
+# -------------------------------------------------------------
+# MAIN FUNCTION: train_model
+# -------------------------------------------------------------
+def train_model(model, loaders, epochs, lr, device, patience=3):
+
+    # 1️⃣ Define loss, optimizer, and learning rate scheduler
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode="max", factor=0.5, patience=1
+    )
+
+    writer = SummaryWriter("runs/recycle") # For TensorBoard visualization
+
+    # Print training setup summary
+    pstats = _num_params(model)
+    classes_line = ", ".join(_class_counts(loaders["train"].dataset)) or "(classes unavailable)"
+    
+    print("\n" + "=" * 68)
+    print(" TRAINING START ")
+    print("-" * 68)
+    print(f" epochs: {epochs} | lr: {lr:.2e} | patience: {patience} | bs: {loaders['train'].batch_size}")
+    print(f" train batches: {len(loaders['train'])} | val batches: {len(loaders['val'])}")
+    print(f" params: {pstats['trainable']:,} trainable / {pstats['total']:,} total")
+    print(f" classes: {classes_line}")
+    print("=" * 68 + "\n")
+
+    # 2️⃣ Initialize trackers
+    best_val_acc, best_state, bad = 0.0, None, 0
+    history = {"train_loss": [], "val_loss": [], "train_acc": [], "val_acc": []}
+    saved_final = False
+
+    # 3️⃣ Training loop over all epochs
+    for ep in range(1, epochs + 1):
+        ep_t0 = time.time()
+
+        # Run training + validation
+        tr_loss, tr_acc = run_epoch(model, loaders["train"], criterion, optimizer, device, True,  ep, epochs)
+        va_loss, va_acc = run_epoch(model, loaders["val"],   criterion, optimizer, device, False, ep, epochs)
+
+        # Update scheduler based on validation accuracy
+        scheduler.step(va_acc)
+
+        # Store metrics in history
+        history["train_loss"].append(tr_loss); history["val_loss"].append(va_loss)
+        history["train_acc"].append(tr_acc);   history["val_acc"].append(va_acc)
+
+        # Log to TensorBoard
+        writer.add_scalar("loss/train", tr_loss, ep)
+        writer.add_scalar("loss/val",   va_loss, ep)
+        writer.add_scalar("acc/train",  tr_acc,  ep)
+        writer.add_scalar("acc/val",    va_acc,  ep)
+        writer.flush()
+
+        
+        # Check for improvement
+        improved = va_acc > best_val_acc
+        if improved:
+            best_val_acc, best_state, bad = va_acc, model.state_dict(), 0
+            Path("models").mkdir(parents=True, exist_ok=True)
+            torch.save(best_state, "models/resnet18_best.pt")
+        else:
+            bad += 1  # No improvement counter
+
+        # Print summary for this epoch
+        print("\n" + "-" * 68)
+        print(f" EPOCH {ep:02d}/{epochs:02d}  |  time { _fmt_time(time.time()-ep_t0) }  |  lr { _current_lrs(optimizer) }")
+        print("-" * 68)
+        print(f"   Train  •  loss {tr_loss:.4f}  |  acc {tr_acc:.4f}")
+        print(f"   Val    •  loss {va_loss:.4f}  |  acc {va_acc:.4f}  |  " f"{'NEW BEST ✓' if improved else f'best {best_val_acc:.4f}'}")
+        
+        if improved:
+            print("   Saved  •  models/resnet18_best.pt")
+
+        # Early stopping
+        if bad >= patience:
+            print("-" * 68)
+            print(f" EARLY STOP  •  no val acc improvement for {patience} epoch(s)")
+            print("-" * 68)
+            _save_final_figure(history, best_val_acc, outpath="reports/curves_final.png", title="ResNet (Fine-tune)")
+            saved_final = True
+            break
+
+    # 4️⃣ Load the best model weights
+    if best_state is not None:
+        model.load_state_dict(best_state)
+
+    # 5️⃣ Save final results if not done during early stop
+    if not saved_final:
+        _save_final_figure(history, best_val_acc, outpath="reports/curves_final.png", title="ResNet (Fine-tune)")
+
+    writer.close()
+    
+    # 6️⃣ Summary printout
+    print("\n" + "=" * 68)
+    print(f" TRAINING DONE  •  best val acc = {best_val_acc:.4f}")
+    print(" Saved final curves: reports/curves_final.png")
+    print("=" * 68 + "\n")
+
+    return model, best_val_acc