added model descriptions

app.py

@@ -7,49 +7,37 @@ yolov8_result = os.path.join(os.getcwd(), "data/xai/yolov8.png")
 yolov5_dff = os.path.join(os.getcwd(), "data/xai/yolov5_dff.png")
 yolov8_dff = os.path.join(os.getcwd(), "data/xai/yolov8_dff.png")
 
+architecture_description_yolov5 = """
+# YOLOv5 Architecture
+
+- **Backbone**: Uses **CSPDarknet53** for feature extraction with **ResNet**-like residual connections.
+- **Neck**: **PANet** and **FPN** aggregate features at multiple scales.
+- **Head**: Predicts bounding boxes (x, y, w, h), class probabilities, and objectness scores.
+- **Loss Functions**: **CIoU** for bounding box regression, **cross-entropy** for classification.
+- **Grid-based Detection**: Divides input into grid cells predicting multiple bounding boxes.
+- **Non-Maximum Suppression (NMS)**: Filters overlapping boxes with high confidence.
+- **Customizable**: Hyperparameters for fine-tuning on custom datasets.
+- **Real-Time Detection**: Optimized for fast, accurate predictions.
+- **Multi-Scale**: Effective for detecting both small and large objects.
+- **Pre-trained Models**: Available for common tasks like detecting people, vehicles, etc.
+"""
 
-architecture_description = """
-# YOLOv5 Architecture Overview
-
-YOLOv5 consists of three main components: the **backbone**, the **neck**, and the **head**.
-
-### 1. **Backbone**: Feature Extraction
-- **CSPDarknet53**: A modified version of **Darknet53**, leveraging **CSPNet** for improved gradient flow and memory usage.
-- **Residual Connections**: Utilizes **ResNet**-like residual connections to enable deeper learning without vanishing gradients.
-- **Focus Layer**: Performs convolutional downsampling to focus on key image features before feature extraction.
-
-### 2. **Neck**: Aggregation of Features
-- **PANet (Path Aggregation Network)**: Used for better feature aggregation and enhanced information flow across scales.
-- **FPN (Feature Pyramid Network)**: Helps in detecting objects at multiple scales by generating a pyramid of feature maps.
-- **Upsample and Downsample**: Combines low-level and high-level features for accurate localization and detection.
-
-### 3. **Head**: Detection and Output Generation
-- **Bounding Box Prediction**: Predicts **center (x, y)**, **width (w)**, and **height (h)** for each bounding box.
-- **Class Prediction**: Outputs the **class probabilities** for each detected object.
-- **Objectness Score**: Predicts whether the bounding box contains an object.
-- **Anchor Boxes**: Uses predefined anchor boxes to assist with bounding box prediction.
-
-### 4. **Detection Layer**: Grid-Based Prediction
-- **Grid-based Prediction**: Divides the image into a grid where each cell predicts multiple bounding boxes.
-- **Non-Maximum Suppression (NMS)**: Filters out redundant bounding boxes with high overlap based on confidence scores.
-
-### 5. **Loss Function**:
-- **CIoU Loss**: Used for bounding box regression to measure the overlap, aspect ratio, and center distance.
-- **Binary Cross-Entropy**: For **objectness score** prediction.
-- **Cross-Entropy**: For **classification** of the detected objects.
-
-### 6. **Post-Processing**:
-- **Non-Maximum Suppression (NMS)**: Eliminates duplicate boxes with high overlap, keeping only the most confident predictions.
-
-### Summary:
-1. **Input Image**: Pre-processed and fed into the model.
-2. **Backbone**: Feature extraction with **CSPDarknet53** and **Residual Connections**.
-3. **Neck**: Feature aggregation with **PANet** and **FPN**.
-4. **Head**: Outputs bounding boxes, objectness scores, and class probabilities.
-5. **Loss Functions**: CIoU, binary cross-entropy, and cross-entropy.
-6. **Post-Processing**: **NMS** to filter overlapping detections.
+architecture_description_yolov8s = """
+# YOLOv8s Architecture
+
+- **Backbone**: Uses **CSPDarknet** with efficient feature extraction layers.
+- **Neck**: Incorporates **FPN** and **PANet** for multi-scale feature aggregation.
+- **Head**: Predicts bounding boxes, objectness scores, and class probabilities.
+- **Loss Functions**: **CIoU** for bounding box regression, **cross-entropy** for classification.
+- **Grid-based Prediction**: Outputs predictions per grid cell using multiple anchor boxes.
+- **Post-Processing**: **NMS** removes redundant boxes with high overlap.
+- **Efficient**: Optimized for real-time, low-latency inference.
+- **Multi-Task**: Supports tasks like classification, detection, and segmentation.
+- **High Accuracy**: Achieves state-of-the-art results on small and medium datasets.
+- **Lightweight**: Efficient architecture for smaller models with reduced computational cost.
 """
 
+
 # Netron HTML templates
 def get_netron_html(model_url):
     return f"""
@@ -103,13 +91,14 @@ with gr.Blocks(css=custom_css) as demo:
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown(architecture_description)
+            gr.Markdown(architecture_description_yolov5)
             gr.HTML(get_netron_html(yolov5_url))
             gr.Image(yolov5_result, label="Detections & Interpretability Map")
             gr.Image(yolov5_dff, label="Feature Factorization & discovered concept")
             
 
         with gr.Column():
+            gr.Markdown(architecture_description_yolov8)
             gr.HTML(get_netron_html(yolov8_url))
             gr.Image(yolov8_result, label="Detections & Interpretability Map")
             gr.Image(yolov8_dff, label="Feature Factorization & discovered concept")