Update app.py

app.py

@@ -9,12 +9,16 @@ from torchvision import datasets, transforms
 import numpy as np
 import threading
 
-app = FastAPI(title="3D CNN Visualizer + MNIST", version="0.2.0")
+app = FastAPI(title="3D CNN Visualizer + MNIST + Online Learning", version="0.3.0")
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 MODEL = None
+OPTIMIZER = None
+CRITERION = None
 TRAINING_DONE = False
 TRAINING_ERROR = None
+MODEL_LOCK = threading.Lock()
 
 
 class SimpleCNN(nn.Module):
@@ -29,8 +33,8 @@ class SimpleCNN(nn.Module):
     """
     def __init__(self):
         super().__init__()
-        self.conv1 = nn.Conv2d(1, 4, kernel_size=5)   # 28 -> 24
-        self.conv2 = nn.Conv2d(4, 8, kernel_size=5)   # 12 -> 8
+        self.conv1 = nn.Conv2d(1, 4, kernel_size=5)
+        self.conv2 = nn.Conv2d(4, 8, kernel_size=5)
         self.pool = nn.MaxPool2d(2, 2)
         self.fc = nn.Linear(8 * 4 * 4, 10)
 
@@ -59,14 +63,15 @@ class SimpleCNN(nn.Module):
         return x
 
 
-def train_model():
-    global MODEL, TRAINING_DONE, TRAINING_ERROR
+def train_model_initial():
+    """Offline initial training on MNIST subset."""
+    global MODEL, OPTIMIZER, CRITERION, TRAINING_DONE, TRAINING_ERROR
     try:
         transform = transforms.ToTensor()
         train_dataset = datasets.MNIST(
             root="./data", train=True, download=True, transform=transform
         )
-        subset_size = min(10000, len(train_dataset))
+        subset_size = min(20000, len(train_dataset))
         train_subset = torch.utils.data.Subset(train_dataset, list(range(subset_size)))
         loader = DataLoader(train_subset, batch_size=128, shuffle=True)
 
@@ -75,7 +80,7 @@ def train_model():
         criterion = nn.CrossEntropyLoss()
 
         model.train()
-        epochs = 1
+        epochs = 2
         for _ in range(epochs):
             for images, labels in loader:
                 images, labels = images.to(DEVICE), labels.to(DEVICE)
@@ -86,21 +91,29 @@ def train_model():
                 optimizer.step()
 
         model.eval()
-        MODEL = model
+        with MODEL_LOCK:
+            MODEL = model
+            OPTIMIZER = optimizer
+            CRITERION = criterion
         TRAINING_DONE = True
     except Exception as e:
         TRAINING_ERROR = str(e)
         TRAINING_DONE = False
 
 
-# 학습을 백그라운드에서 시작
-threading.Thread(target=train_model, daemon=True).start()
+# start initial training in background
+threading.Thread(target=train_model_initial, daemon=True).start()
 
 
 class PredictRequest(BaseModel):
     pixels: list[float]  # 28*28 = 784
 
 
+class FeedbackRequest(BaseModel):
+    pixels: list[float]  # 28*28
+    label: int           # 0~9
+
+
 @app.get("/", response_class=HTMLResponse)
 async def index():
     return HTML_PAGE
@@ -115,6 +128,26 @@ async def status():
     }
 
 
+def _forward_with_activations(x_tensor: torch.Tensor):
+    """Run model forward and return probs + activations as python lists."""
+    logits, acts = MODEL(x_tensor, return_activations=True)
+    probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
+
+    conv1 = acts["conv1"].detach().cpu().numpy()[0].tolist()   # [4,24,24]
+    pool1 = acts["pool1"].detach().cpu().numpy()[0].tolist()   # [4,12,12]
+    conv2 = acts["conv2"].detach().cpu().numpy()[0].tolist()   # [8,8,8]
+    pool2 = acts["pool2"].detach().cpu().numpy()[0].tolist()   # [8,4,4]
+    flat = acts["flat"].detach().cpu().numpy()[0].tolist()     # [128]
+
+    return probs, {
+        "conv1": conv1,
+        "pool1": pool1,
+        "conv2": conv2,
+        "pool2": pool2,
+        "flat": flat,
+    }
+
+
 @app.post("/predict")
 async def predict(req: PredictRequest):
     if TRAINING_ERROR:
@@ -136,28 +169,74 @@ async def predict(req: PredictRequest):
 
     arr = np.array(req.pixels, dtype=np.float32).reshape(1, 1, 28, 28)
     x = torch.from_numpy(arr).to(DEVICE)
-    with torch.no_grad():
-        logits, acts = MODEL(x, return_activations=True)
-        probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
 
-        conv1 = acts["conv1"].cpu().numpy()[0].tolist()   # [4,24,24]
-        pool1 = acts["pool1"].cpu().numpy()[0].tolist()   # [4,12,12]
-        conv2 = acts["conv2"].cpu().numpy()[0].tolist()   # [8,8,8]
-        pool2 = acts["pool2"].cpu().numpy()[0].tolist()   # [8,4,4]
-        flat = acts["flat"].cpu().numpy()[0].tolist()     # [128]
+    with MODEL_LOCK:
+        MODEL.eval()
+        with torch.no_grad():
+            probs, acts = _forward_with_activations(x)
 
     predicted_class = int(probs.argmax())
 
     return {
         "predicted_class": predicted_class,
         "probabilities": probs.tolist(),
-        "activations": {
-            "conv1": conv1,
-            "pool1": pool1,
-            "conv2": conv2,
-            "pool2": pool2,
-            "flat": flat,
-        },
+        "activations": acts,
+    }
+
+
+@app.post("/feedback")
+async def feedback(req: FeedbackRequest):
+    """Online learning: single-sample update with user-labeled digit."""
+    if TRAINING_ERROR:
+        return JSONResponse(
+            status_code=500,
+            content={"error": "Training failed", "detail": TRAINING_ERROR},
+        )
+    if not TRAINING_DONE or MODEL is None or OPTIMIZER is None or CRITERION is None:
+        return JSONResponse(
+            status_code=503,
+            content={"error": "Model not ready yet."},
+        )
+
+    if len(req.pixels) != 28 * 28:
+        return JSONResponse(
+            status_code=400,
+            content={"error": "pixels must have length 784 (28x28)"},
+        )
+    if not (0 <= req.label <= 9):
+        return JSONResponse(
+            status_code=400,
+            content={"error": "label must be between 0 and 9"},
+        )
+
+    arr = np.array(req.pixels, dtype=np.float32).reshape(1, 1, 28, 28)
+    x = torch.from_numpy(arr).to(DEVICE)
+    y = torch.tensor([req.label], dtype=torch.long, device=DEVICE)
+
+    with MODEL_LOCK:
+        MODEL.train()
+        loss_val = None
+        # do a few small gradient steps on this sample
+        for _ in range(3):
+            OPTIMIZER.zero_grad()
+            logits = MODEL(x)
+            loss = CRITERION(logits, y)
+            loss.backward()
+            OPTIMIZER.step()
+            loss_val = float(loss.item())
+        MODEL.eval()
+        # re-run forward to get updated prediction & activations
+        with torch.no_grad():
+            probs, acts = _forward_with_activations(x)
+
+    predicted_class = int(probs.argmax())
+
+    return {
+        "status": "ok",
+        "loss": loss_val,
+        "predicted_class": predicted_class,
+        "probabilities": probs.tolist(),
+        "activations": acts,
     }
 
 
@@ -544,6 +623,7 @@ HTML_PAGE = r"""
             const [statusText, setStatusText] = useState("Checking model status...");
             const [lastPrediction, setLastPrediction] = useState(null);
             const [padKey, setPadKey] = useState(0);
+            const [feedbackBusy, setFeedbackBusy] = useState(false);
 
             useEffect(() => {
                 audio.init();
@@ -577,13 +657,17 @@ HTML_PAGE = r"""
 
             const delay = ms => new Promise(r => setTimeout(r, ms));
 
+            const flattenInput = (grid) => {
+                const flat = [];
+                for(let y=0; y<28; y++) for(let x=0; x<28; x++) flat.push(grid[y][x]);
+                return flat;
+            };
+
             const run = async () => {
-                if(processing) return;
+                if(processing || feedbackBusy) return;
                 setProcessing(true);
 
-                // 1) 입력 플래튼
-                const flat = [];
-                for(let y=0; y<28; y++) for(let x=0; x<28; x++) flat.push(inputData[y][x]);
+                const flat = flattenInput(inputData);
 
                 let probs = null;
                 let predClass = null;
@@ -601,7 +685,8 @@ HTML_PAGE = r"""
                         acts = json.activations;
                         setLastPrediction({
                             cls: predClass,
-                            conf: probs[predClass]
+                            conf: probs[predClass],
+                            probs,
                         });
                     } else {
                         alert("Error: " + (json.error || "Unknown"));
@@ -615,15 +700,13 @@ HTML_PAGE = r"""
                     return;
                 }
 
-                // 2) PyTorch activations → 시각화용 포맷 변환
-                // conv1 / pool1 / conv2 / pool2 는 그대로 사용 (depth x h x w)
+                // convert activations for visualization
                 const conv1 = acts.conv1 || [];
                 const pool1 = acts.pool1 || [];
                 const conv2 = acts.conv2 || [];
                 const pool2 = acts.pool2 || [];
-                const flatVec = acts.flat || [];           // length 128
+                const flatVec = acts.flat || [];
 
-                // flat: 1 x 8 x 16으로 reshape
                 const flatGrid = [];
                 for(let y=0; y<8; y++) {
                     const row = [];
@@ -633,38 +716,105 @@ HTML_PAGE = r"""
                     }
                     flatGrid.push(row);
                 }
-                const flatData = [flatGrid];  // depth=1
-
-                // fc: [probability] 형태로 감싸서 시각화
+                const flatData = [flatGrid];
                 const fcData = [probs.map(p => [p])];
 
-                // 3) 단계별로 activations 상태 업데이트 (애니메이션)
                 setActivations(prev => ({...prev, input: [inputData]}));
-                setStep(0); audio.playStep(0); await delay(400);
+                setStep(0); audio.playStep(0); await delay(250);
 
                 setActivations(prev => ({...prev, conv1}));
-                setStep(1); audio.playStep(1); await delay(400);
+                setStep(1); audio.playStep(1); await delay(250);
 
                 setActivations(prev => ({...prev, pool1}));
-                setStep(2); audio.playStep(2); await delay(400);
+                setStep(2); audio.playStep(2); await delay(250);
 
                 setActivations(prev => ({...prev, conv2}));
-                setStep(3); audio.playStep(3); await delay(400);
+                setStep(3); audio.playStep(3); await delay(250);
 
                 setActivations(prev => ({...prev, pool2}));
-                setStep(4); audio.playStep(4); await delay(400);
+                setStep(4); audio.playStep(4); await delay(250);
 
                 setActivations(prev => ({...prev, flat: flatData}));
-                setStep(5); audio.playStep(5); await delay(400);
+                setStep(5); audio.playStep(5); await delay(250);
 
                 setActivations(prev => ({...prev, fc: fcData}));
                 setStep(6); audio.playStep(6);
                 
-                await delay(1500);
+                await delay(1000);
                 setProcessing(false);
                 setStep(-1);
             };
 
+            const sendFeedback = async (correctLabel) => {
+                if(!lastPrediction) return;
+                if(processing || feedbackBusy) return;
+                setFeedbackBusy(true);
+                setStatusText(`Online update with label ${correctLabel} ...`);
+
+                const flat = flattenInput(inputData);
+
+                try {
+                    const res = await fetch("/feedback", {
+                        method: "POST",
+                        headers: { "Content-Type": "application/json" },
+                        body: JSON.stringify({ pixels: flat, label: correctLabel }),
+                    });
+                    const json = await res.json();
+                    if(!res.ok) {
+                        alert("Feedback error: " + (json.error || "Unknown"));
+                        setFeedbackBusy(false);
+                        return;
+                    }
+
+                    const probs = json.probabilities;
+                    const acts = json.activations;
+                    const predClass = json.predicted_class;
+
+                    setLastPrediction({
+                        cls: predClass,
+                        conf: probs[predClass],
+                        probs,
+                    });
+
+                    // update visualization with new activations
+                    const conv1 = acts.conv1 || [];
+                    const pool1 = acts.pool1 || [];
+                    const conv2 = acts.conv2 || [];
+                    const pool2 = acts.pool2 || [];
+                    const flatVec = acts.flat || [];
+
+                    const flatGrid = [];
+                    for(let y=0; y<8; y++) {
+                        const row = [];
+                        for(let x=0; x<16; x++) {
+                            const idx = y * 16 + x;
+                            row.push(flatVec[idx] || 0);
+                        }
+                        flatGrid.push(row);
+                    }
+                    const flatData = [flatGrid];
+                    const fcData = [probs.map(p => [p])];
+
+                    setActivations(prev => ({
+                        ...prev,
+                        input: [inputData],
+                        conv1,
+                        pool1,
+                        conv2,
+                        pool2,
+                        flat: flatData,
+                        fc: fcData,
+                    }));
+
+                    setStatusText(`Online updated with label ${correctLabel} (loss ~${json.loss.toFixed(4)})`);
+                } catch (e) {
+                    console.error(e);
+                    alert("Feedback request failed.");
+                } finally {
+                    setFeedbackBusy(false);
+                }
+            };
+
             return (
                 <div className="w-full h-screen relative bg-black font-mono">
                     <Canvas shadows camera={{ position: [25, 10, 5], fov: 45 }}>
@@ -692,7 +842,7 @@ HTML_PAGE = r"""
                                     <Cpu className="text-neon-green animate-pulse" /> DEEP <span className="text-neon-green">CNN</span>
                                 </h1>
                                 <div className="text-xs text-green-400 mt-2 flex items-center gap-2">
-                                    <Activity size={12} /> {processing ? "PROCESSING TENSORS..." : "ONLINE"}
+                                    <Activity size={12} /> {(processing || feedbackBusy) ? "PROCESSING / UPDATING..." : "ONLINE"}
                                 </div>
                                 <div className="text-[10px] text-green-500 mt-1">{statusText}</div>
                                 {lastPrediction && (
@@ -715,15 +865,43 @@ HTML_PAGE = r"""
                                     <div className="text-xs font-bold tracking-widest flex items-center gap-2">
                                         <Scan size={14} /> INPUT SENSOR
                                     </div>
-                                    <button onClick={reset} disabled={processing} className="hover:text-white transition-colors">
+                                    <button onClick={reset} disabled={processing || feedbackBusy} className="hover:text-white transition-colors">
                                         <RotateCcw size={16} />
                                     </button>
                                 </div>
-                                <DrawingPad key={padKey} data={inputData} onChange={setInputData} disabled={processing} />
-                                <button onClick={run} disabled={processing} className="w-full mt-4 py-3 rounded btn-holo flex justify-center items-center gap-2 font-bold transition-all">
-                                    {processing ? <Activity className="animate-spin" size={18} /> : <Play size={18} fill="currentColor" />}
-                                    {processing ? 'CALCULATING...' : 'RUN INFERENCE'}
+                                <DrawingPad key={padKey} data={inputData} onChange={setInputData} disabled={processing || feedbackBusy} />
+                                <button onClick={run} disabled={processing || feedbackBusy} className="w-full mt-4 py-3 rounded btn-holo flex justify-center items-center gap-2 font-bold transition-all">
+                                    {(processing || feedbackBusy) ? <Activity className="animate-spin" size={18} /> : <Play size={18} fill="currentColor" />}
+                                    {(processing || feedbackBusy) ? 'RUNNING / UPDATING...' : 'RUN INFERENCE'}
                                 </button>
+
+                                {lastPrediction && (
+                                    <div className="mt-4">
+                                        <div className="text-[11px] text-green-400 mb-2">
+                                            Select the <span className="font-bold text-neon-green">correct digit</span> for online learning:
+                                        </div>
+                                        <div className="grid grid-cols-5 gap-2">
+                                            {Array.from({length:10}).map((_, d) => {
+                                                const isPred = lastPrediction.cls === d;
+                                                return (
+                                                    <button
+                                                        key={d}
+                                                        onClick={() => sendFeedback(d)}
+                                                        disabled={feedbackBusy || processing}
+                                                        className={
+                                                            "border rounded px-2 py-1 text-xs " +
+                                                            (isPred 
+                                                                ? "border-neon-green text-neon-green bg-black/40" 
+                                                                : "border-green-800 text-green-400 hover:border-neon-green hover:text-neon-green")
+                                                        }
+                                                    >
+                                                        {d}
+                                                    </button>
+                                                );
+                                            })}
+                                        </div>
+                                    </div>
+                                )}
                             </div>
 
                             <div className="hud-panel p-5 hidden md:block rounded-t-xl min-w-[260px] border-b-0">