Update app.py

app.py

@@ -1,20 +1,10 @@
-import os
 import torch
 import torchvision.transforms as transforms
-import torchvision
 import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
-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")
+import os
 
 # === Simple CNN Model Definition ===
 class SimpleCNN(nn.Module):
@@ -33,33 +23,28 @@ class SimpleCNN(nn.Module):
         x = F.relu(self.fc1(x))
         return self.fc2(x)
 
-# === Load DCLR model (for inference tab) ===
+# === Model Loading ===
 model = SimpleCNN()
-if os.path.exists(DCLR_MODEL_PATH):
-    model.load_state_dict(torch.load(DCLR_MODEL_PATH, map_location=torch.device('cpu')))
+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')))
     model.eval()
-    print(f"Model loaded successfully from {DCLR_MODEL_PATH}")
+    print(f"Model loaded successfully from {model_path}")
 else:
-    print(f"Warning: Model file '{DCLR_MODEL_PATH}' not found. Run train_dclr_model.py.")
+    print(f"Warning: Model file '{model_path}' not found. Please run train_dclr_model.py first.")
 
 # === CIFAR-10 Class Labels ===
 class_labels = ['plane','car','bird','cat','deer','dog','frog','horse','ship','truck']
 
-# === Image Preprocessing (consistent with training normalization) ===
+# === Image Preprocessing ===
 preprocess = transforms.Compose([
     transforms.Resize(32),
     transforms.ToTensor(),
-    transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
+    transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
 ])
 
-# === CIFAR-10 Test Loader for Benchmark Mode ===
-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) ===
+# === Inference Function ===
 def inference(input_image: Image.Image):
     if model.training:
         model.eval()
@@ -70,106 +55,39 @@ 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 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)
-    class_total = np.zeros(10)
-
-    with torch.no_grad():
-        for inputs, labels in test_loader:
-            outputs = local_model(inputs)
-            _, predicted = outputs.max(1)
-            total += labels.size(0)
-            correct += predicted.eq(labels).sum().item()
-            c = (predicted == labels).squeeze()
-            for i in range(len(labels)):
-                label = labels[i].item()
-                class_correct[label] += c[i].item()
-                class_total[label] += 1
-
-    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)}
-
-    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
+# === Results Viewer Function ===
+def show_results(input_image: Image.Image):
+    preds = inference(input_image)
 
-# === 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."
+    # 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
 
-# === Prepare CIFAR-10 Sample Gallery (one per class with captions) ===
-sample_dir = "examples"
-os.makedirs(sample_dir, exist_ok=True)
+    # Load final test accuracy number
+    test_acc_text = "Final test accuracy not available."
+    if os.path.exists("final_test_accuracy.txt"):
+        with open("final_test_accuracy.txt", "r") as f:
+            test_acc_value = f.read().strip()
+        test_acc_text = f"Final Test Accuracy: {test_acc_value}%"
 
-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
+    return preds, perf_plot, acc_plot, test_acc_text
 
 # === Gradio Interface Setup ===
-with gr.Blocks() as demo:
-    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 (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)
-
-        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])
+example_images = []
 
-    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)
+interface = gr.Interface(
+    fn=show_results,
+    inputs=gr.Image(type='pil', label='Upload Image'),
+    outputs=[
+        gr.Label(num_top_classes=3, label='Predictions'),
+        gr.Image(type='filepath', label='Training Performance'),
+        gr.Image(type='filepath', label='Final Test Accuracy Plot'),
+        gr.Textbox(label='Final Test Accuracy')
+    ],
+    title='CIFAR-10 Image Classification with DCLR Optimizer',
+    description='Upload an image to see predictions. Training/test plots and accuracy show benchmark results on CIFAR-10.',
+    examples=example_images
+)
 
 if __name__ == '__main__':
-    demo.launch()
+    interface.launch()