Update app.py

app.py

@@ -24,17 +24,51 @@ try:
     weights_loaded = True                                                                                                                                                                                            
     print("Weights loaded successfully!")                                                                                                                                                                            
 except FileNotFoundError:                                                                                                                                                                                            
-    print("Warning: model_weights.npz not found. App will run in warning mode.")                                                                                                                                     
+    print("Warning: model_weights.npz not found.")                                                                                                                                                                   
 except Exception as e:                                                                                                                                                                                               
     print(f"Warning: error loading weights: {e}")                                                                                                                                                                    
                                                                                                                                                                                                                      
+# Helper to stitch feature maps into a nice grid                                                                                                                                                                     
+def make_grid(feature_maps, cols=4):                                                                                                                                                                                 
+    n, h, w = feature_maps.shape                                                                                                                                                                                     
+    rows = (n + cols - 1) // cols                                                                                                                                                                                    
+                                                                                                                                                                                                                     
+    # Pad borders between filters so they look separate                                                                                                                                                              
+    padding = 2                                                                                                                                                                                                      
+    grid_h = rows * h + (rows - 1) * padding                                                                                                                                                                         
+    grid_w = cols * w + (cols - 1) * padding                                                                                                                                                                         
+    grid = np.zeros((grid_h, grid_w), dtype=np.uint8)                                                                                                                                                                
+                                                                                                                                                                                                                     
+    for idx in range(n):                                                                                                                                                                                             
+        r = idx // cols                                                                                                                                                                                              
+        c = idx % cols                                                                                                                                                                                               
+                                                                                                                                                                                                                     
+        f_map = feature_maps[idx]                                                                                                                                                                                    
+        f_min, f_max = f_map.min(), f_map.max()                                                                                                                                                                      
+        # Normalize to [0, 255] for image display                                                                                                                                                                    
+        if f_max > f_min:                                                                                                                                                                                            
+            f_map = 255.0 * (f_map - f_min) / (f_max - f_min)                                                                                                                                                        
+        else:                                                                                                                                                                                                        
+            f_map = f_map * 0                                                                                                                                                                                        
+                                                                                                                                                                                                                     
+        f_map = f_map.astype(np.uint8)                                                                                                                                                                               
+                                                                                                                                                                                                                     
+        y_start = r * (h + padding)                                                                                                                                                                                  
+        x_start = c * (w + padding)                                                                                                                                                                                  
+        grid[y_start:y_start+h, x_start:x_start+w] = f_map                                                                                                                                                           
+                                                                                                                                                                                                                     
+    img = Image.fromarray(grid)                                                                                                                                                                                      
+    # Scale up using nearest-neighbor to keep pixels clean and sharp                                                                                                                                                 
+    img = img.resize((grid_w * 12, grid_h * 12), Image.Resampling.NEAREST)                                                                                                                                           
+    return img                                                                                                                                                                                                       
+                                                                                                                                                                                                                     
 # 3. Predict function                                                                                                                                                                                                
 def predict(input_image):                                                                                                                                                                                            
     if not weights_loaded:                                                                                                                                                                                           
-        return {"Error: please upload 'model_weights.npz' to your Space files": 1.0}                                                                                                                                 
+        return {"Error: please upload 'model_weights.npz'": 1.0}, None, None                                                                                                                                         
                                                                                                                                                                                                                      
     if input_image is None:                                                                                                                                                                                          
-        return "No image drawn"                                                                                                                                                                                      
+        return "No image drawn", None, None                                                                                                                                                                          
                                                                                                                                                                                                                      
     if isinstance(input_image, dict):                                                                                                                                                                                
         img = input_image['composite']                                                                                                                                                                               
@@ -82,31 +116,40 @@ def predict(input_image):
         canvas.paste(resized, (offset_x, offset_y))                                                                                                                                                                  
         arr = np.array(canvas)                                                                                                                                                                                       
                                                                                                                                                                                                                      
-    # Normalize                                                                                                                                                                                                      
     x = (arr / 255.0) - 0.5                                                                                                                                                                                          
     x = x[np.newaxis, :, :] # (1, 28, 28)                                                                                                                                                                            
                                                                                                                                                                                                                      
-    # Run forward pass                                                                                                                                                                                               
-    out1 = conv.forward(x)                                                                                                                                                                                           
-    out2 = gelu.forward(out1)                                                                                                                                                                                        
-    out3 = pool.forward(out2)                                                                                                                                                                                        
-    out4 = flatten.forward(out3[np.newaxis, :, :, :])                                                                                                                                                                
-    logits = dense.forward(out4)[0]                                                                                                                                                                                  
+    # Forward pass and record intermediate activations                                                                                                                                                               
+    out_conv = conv.forward(x)                                                                                                                                                                                       
+    out_gelu = gelu.forward(out_conv)                                                                                                                                                                                
+    out_pool = pool.forward(out_gelu)                                                                                                                                                                                
+                                                                                                                                                                                                                     
+    out_flat = flatten.forward(out_pool[np.newaxis, :, :, :])                                                                                                                                                        
+    logits = dense.forward(out_flat)[0]                                                                                                                                                                              
                                                                                                                                                                                                                      
     # Softmax probabilities                                                                                                                                                                                          
     probs = np.exp(logits - np.max(logits))                                                                                                                                                                          
     probs /= np.sum(probs)                                                                                                                                                                                           
                                                                                                                                                                                                                      
-    # Return dictionary of classes and probabilities                                                                                                                                                                 
-    return {str(i): float(probs[i]) for i in range(10)}                                                                                                                                                              
+    # Generate feature map grids for visualization                                                                                                                                                                   
+    conv_grid = make_grid(out_conv)                                                                                                                                                                                  
+    pool_grid = make_grid(out_pool)                                                                                                                                                                                  
+                                                                                                                                                                                                                     
+    # Return dictionary of classes, plus the two grids                                                                                                                                                               
+    class_probs = {str(i): float(probs[i]) for i in range(10)}                                                                                                                                                       
+    return class_probs, conv_grid, pool_grid                                                                                                                                                                         
                                                                                                                                                                                                                      
 # Gradio Interface layout                                                                                                                                                                                            
 demo = gr.Interface(                                                                                                                                                                                                 
     fn=predict,                                                                                                                                                                                                      
     inputs=gr.Sketchpad(type="pil", image_mode="L"),                                                                                                                                                                 
-    outputs=gr.Label(num_top_classes=3),                                                                                                                                                                             
+    outputs=[                                                                                                                                                                                                        
+        gr.Label(num_top_classes=3, label="Prediction"),                                                                                                                                                             
+        gr.Image(label="Layer 1: Convolutional Activations (12 Filters)", type="pil"),                                                                                                                               
+        gr.Image(label="Layer 2: Max Pooling Outputs (Downsampled Features)", type="pil")                                                                                                                            
+    ],                                                                                                                                                                                                               
     title="CNN from Scratch - Digit Classifier",                                                                                                                                                                     
-    description="Draw a digit in the box to test our scratch-built NumPy model!"                                                                                                                                     
+    description="Draw a digit in the box to predict its value and see inside the model's 'brain' in real-time!"                                                                                                      
 )                                                                                                                                                                                                                    
                                                                                                                                                                                                                      
 if __name__ == "__main__":