Update data_viz.py

data_viz.py

@@ -99,19 +99,21 @@ def data_visualization(model_choice):
 
     elif model_choice == "VGG16":
         st.markdown("""
-    #### 📦 VGG16 Model Summary
-
-    VGG16 is a deep convolutional neural network known for its uniform architecture of `3x3` conv filters and `2x2` max-pooling layers. It ends with 3 fully connected layers.
-
-    Since it's quite large (138 million parameters), we’re showing a high-level overview here:
-
-    | Component            | Description                        |
-    |---------------------|------------------------------------|
-    | Convolutional Blocks| 13 Conv layers + ReLU + MaxPooling |
-    | Fully Connected     | FC1 → FC2 → Output layer (6 units) |
-    | Pretrained Base     | Yes (ImageNet, fine-tuned)         |
-
-    **Total Parameters**: ~`138 million`  
-    **Trainable Parameters**: ~`138 million`  
-    **Estimated Model Size**: ~`500+ MB`  
+    #### 🏗️ VGG16 Architecture (Fine-Tuned)
+    
+    | Layer                         | Input Shape         | Output Shape        | Params       | Trainable |
+    |------------------------------|---------------------|---------------------|--------------|-----------|
+    | Features (13 conv layers)    | [32, 3, 224, 224]   | [32, 512, 7, 7]     | ~7.9M        | ❌        |
+    | AdaptiveAvgPool2d            | [32, 512, 7, 7]     | [32, 512, 7, 7]     | 0            | ❌        |
+    | Linear (fc1)                 | [32, 25088]         | [32, 4096]          | 102,764,544  | ✅        |
+    | Linear (fc2)                 | [32, 4096]          | [32, 1024]          | 4,195,328    | ✅        |
+    | Linear (fc3)                 | [32, 1024]          | [32, 6]             | 6,150        | ✅        |
+    
+    **Total Parameters**: `116,186,502`  
+    **Trainable Parameters**: `106,966,022`  
+    **Non-trainable Parameters**: `9,220,480`  
+    **Estimated Model Size**: `~2.4 GB`
     """)
+    
+    else:
+        st.warning("⚠️ Model summary not available.")