Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import torch
 import torchvision.transforms as transforms
 import torchvision
@@ -5,9 +6,16 @@ import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
-import os
 import numpy as np
 
+# === Paths ===
+ART_DIR = "artifacts"
+DCLR_MODEL_PATH = os.path.join(ART_DIR, "dclr_simple_cnn.pth")
+DCLR_PERF_PNG = os.path.join(ART_DIR, "dclr_training_performance.png")
+DCLR_ACC_PNG = os.path.join(ART_DIR, "dclr_final_test_accuracy.png")
+DCLR_ACC_TXT = os.path.join(ART_DIR, "dclr_final_test_accuracy.txt")
+BENCHMARK_TXT = os.path.join(ART_DIR, "benchmark_results.txt")
+
 # === Simple CNN Model Definition ===
 class SimpleCNN(nn.Module):
     def __init__(self):
@@ -25,29 +33,30 @@ class SimpleCNN(nn.Module):
         x = F.relu(self.fc1(x))
         return self.fc2(x)
 
-# === Model Loading ===
+# === Load DCLR model (for inference tab) ===
 model = SimpleCNN()
-model_path = 'simple_cnn_dclr_tuned.pth'
-
-if os.path.exists(model_path):
-    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+if os.path.exists(DCLR_MODEL_PATH):
+    model.load_state_dict(torch.load(DCLR_MODEL_PATH, map_location=torch.device('cpu')))
     model.eval()
-    print(f"Model loaded successfully from {model_path}")
+    print(f"Model loaded successfully from {DCLR_MODEL_PATH}")
 else:
-    print(f"Warning: Model file '{model_path}' not found. Please run train_dclr_model.py first.")
+    print(f"Warning: Model file '{DCLR_MODEL_PATH}' not found. Run train_dclr_model.py.")
 
 # === CIFAR-10 Class Labels ===
 class_labels = ['plane','car','bird','cat','deer','dog','frog','horse','ship','truck']
 
-# === Image Preprocessing ===
+# === Image Preprocessing (consistent with training normalization) ===
 preprocess = transforms.Compose([
     transforms.Resize(32),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
+    transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
 ])
 
 # === CIFAR-10 Test Loader for Benchmark Mode ===
-test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.ToTensor())
+test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+]))
 test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False)
 
 # === Inference Function (single image) ===
@@ -61,9 +70,16 @@ def inference(input_image: Image.Image):
     confidences = {class_labels[i]: float(probabilities[0,i]) for i in range(len(class_labels))}
     return confidences
 
-# === Benchmark Mode: Evaluate on full test set ===
-def benchmark():
-    model.eval()
+# === Benchmark Mode: Evaluate DCLR on full test set (real-time) ===
+def benchmark_dclr_realtime():
+    if not os.path.exists(DCLR_MODEL_PATH):
+        return "Model missing. Run training first.", {}, None, None
+
+    # Load weights fresh to avoid any accidental state drift
+    local_model = SimpleCNN()
+    local_model.load_state_dict(torch.load(DCLR_MODEL_PATH, map_location=torch.device('cpu')))
+    local_model.eval()
+
     correct = 0
     total = 0
     class_correct = np.zeros(10)
@@ -71,7 +87,7 @@ def benchmark():
 
     with torch.no_grad():
         for inputs, labels in test_loader:
-            outputs = model(inputs)
+            outputs = local_model(inputs)
             _, predicted = outputs.max(1)
             total += labels.size(0)
             correct += predicted.eq(labels).sum().item()
@@ -81,33 +97,79 @@ def benchmark():
                 class_correct[label] += c[i].item()
                 class_total[label] += 1
 
-    overall_acc = 100.0 * correct / total
+    overall_acc = round(100.0 * correct / total, 2)
     classwise_acc = {class_labels[i]: round(100.0 * class_correct[i] / class_total[i], 2) for i in range(10)}
 
-    # Load plots if they exist
-    perf_plot = "training_performance.png" if os.path.exists("training_performance.png") else None
-    acc_plot = "final_test_accuracy.png" if os.path.exists("final_test_accuracy.png") else None
-
-    return overall_acc, classwise_acc, perf_plot, acc_plot
+    perf_plot = DCLR_PERF_PNG if os.path.exists(DCLR_PERF_PNG) else None
+    acc_plot = DCLR_ACC_PNG if os.path.exists(DCLR_ACC_PNG) else None
+
+    return f"{overall_acc}%", classwise_acc, perf_plot, acc_plot
+
+# === Benchmark Comparison: Read real ledger (DCLR vs Adam vs Lion) ===
+def benchmark_comparison():
+    if os.path.exists(BENCHMARK_TXT):
+        with open(BENCHMARK_TXT, "r") as f:
+            return f.read()
+    return "No benchmark_results.txt found. Please run train_dclr_model.py to generate real numbers."
+
+# === Prepare CIFAR-10 Sample Gallery (one per class with captions) ===
+sample_dir = "examples"
+os.makedirs(sample_dir, exist_ok=True)
+
+transform_gallery = transforms.Compose([transforms.ToPILImage()])
+raw_test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.ToTensor())
+
+example_images = []
+seen_classes = set()
+for idx in range(len(raw_test_set)):
+    img, label = raw_test_set[idx]
+    if label not in seen_classes:
+        pil_img = transform_gallery(img)
+        file_path = os.path.join(sample_dir, f"example_{class_labels[label]}.png")
+        pil_img.save(file_path)
+        example_images.append([file_path, f"Sample {class_labels[label]}"])
+        seen_classes.add(label)
+    if len(seen_classes) == 10:
+        break
 
 # === Gradio Interface Setup ===
 with gr.Blocks() as demo:
-    gr.Markdown("## CIFAR-10 Image Classification with DCLR Optimizer")
-    gr.Markdown("Upload an image for prediction, or run Benchmark Mode to see full test accuracy.")
+    gr.Markdown("# DCLR Optimiser — CIFAR-10 Artifact Viewer")
+    gr.Markdown("Upload an image for prediction, or use Benchmark tabs for real test results. All numbers are computed from CIFAR-10 runs and saved as reproducible artifacts.")
 
-    with gr.Tab("Single Image Inference"):
-        inp = gr.Image(type='pil', label='Upload Image')
+    with gr.Tab("Single Image Inference (DCLR)"):
+        inp = gr.Image(type='pil', label='Upload Image (32x32 assumed)')
         out = gr.Label(num_top_classes=3, label='Predictions')
+        perf_img = gr.Image(type='filepath', label='DCLR Training Performance', value=DCLR_PERF_PNG if os.path.exists(DCLR_PERF_PNG) else None)
+        acc_img = gr.Image(type='filepath', label='DCLR Final Test Accuracy Plot', value=DCLR_ACC_PNG if os.path.exists(DCLR_ACC_PNG) else None)
+        acc_text = gr.Textbox(label='DCLR Final Test Accuracy')
+        # If the accuracy text file exists, load it at UI init
+        if os.path.exists(DCLR_ACC_TXT):
+            with open(DCLR_ACC_TXT, "r") as f:
+                acc_text.value = f"Final Test Accuracy: {f.read().strip()}%"
+        # Hook
         inp.change(fn=inference, inputs=inp, outputs=out)
 
-    with gr.Tab("Benchmark Mode"):
-        btn = gr.Button("Run Benchmark on CIFAR-10 Test Set")
-        overall = gr.Textbox(label="Overall Test Accuracy")
-        classwise = gr.JSON(label="Per-Class Accuracy (%)")
-        perf_plot = gr.Image(type='filepath', label='Training Performance')
-        acc_plot = gr.Image(type='filepath', label='Final Test Accuracy Plot')
-
-        btn.click(fn=benchmark, inputs=None, outputs=[overall, classwise, perf_plot, acc_plot])
+        gr.Examples(
+            examples=example_images,
+            inputs=inp,
+            label="CIFAR-10 Samples (one per class)"
+        )
+
+    with gr.Tab("Benchmark Mode (DCLR real-time)"):
+        btn = gr.Button("Run DCLR Benchmark on CIFAR-10 Test Set")
+        overall = gr.Textbox(label="Overall Test Accuracy (DCLR)")
+        classwise = gr.JSON(label="Per-Class Accuracy (%) (DCLR)")
+        perf_plot = gr.Image(type='filepath', label='DCLR Training Performance')
+        acc_plot = gr.Image(type='filepath', label='DCLR Final Test Accuracy Plot')
+
+        btn.click(fn=benchmark_dclr_realtime, inputs=None, outputs=[overall, classwise, perf_plot, acc_plot])
+
+    with gr.Tab("Benchmark Comparison (DCLR vs Adam vs Lion)"):
+        gr.Markdown("Reads real results from artifacts/benchmark_results.txt produced by training.")
+        show_btn = gr.Button("Show Real Benchmark Ledger")
+        ledger_box = gr.Textbox(label="Benchmark Results", lines=10)
+        show_btn.click(fn=benchmark_comparison, inputs=None, outputs=ledger_box)
 
 if __name__ == '__main__':
     demo.launch()