Update app.py

app.py

@@ -4,60 +4,148 @@ import numpy as np
 import pandas as pd
 from collections import Counter
 from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
-
-# # Download YOLOv10 model from Hugging Face
-# MODEL_PATH = hf_hub_download(
-#     repo_id="ibrahim313/Bioengineering_Query_Tool_image_based",
-#     filename="best.pt"
-# )
+import plotly.express as px
+import plotly.graph_objects as go
 
 # Load the model
 model = YOLO("best.pt")
 
-def process_image(image):
+def create_size_distribution_plot(df):
+    """Create a box plot of cell sizes for each class."""
+    fig = px.box(df, x="class_name", y="area", title="Cell Size Distribution by Type")
+    fig.update_layout(
+        xaxis_title="Cell Type",
+        yaxis_title="Area (pixels²)",
+        template="plotly_white"
+    )
+    return fig
+
+def create_density_heatmap(df, image_shape):
+    """Create a heatmap showing cell density."""
+    heatmap = np.zeros(image_shape[:2])
+    for _, row in df.iterrows():
+        center_x = int((row['x_min'] + row['x_max']) / 2)
+        center_y = int((row['y_min'] + row['y_max']) / 2)
+        heatmap[max(0, center_y-20):min(image_shape[0], center_y+20),
+                max(0, center_x-20):min(image_shape[1], center_x+20)] += 1
+    
+    fig = go.Figure(data=go.Heatmap(z=heatmap))
+    fig.update_layout(title="Cell Density Heatmap")
+    return fig
+
+def process_image(image, conf_threshold=0.25):
     """Detect cells in the image, extract attributes, and return results."""
+    if image is None:
+        return None, "No image uploaded", None, None, None
+    
     # Convert image to RGB
     image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     
     # Perform detection
-    results = model.predict(source=image_rgb, imgsz=640, conf=0.25)
+    results = model.predict(source=image_rgb, imgsz=640, conf=conf_threshold)
     
     # Get annotated image
     annotated_img = results[0].plot()
     
     # Extract detection data
     detections = results[0].boxes.data if results[0].boxes is not None else []
+    
     if len(detections) > 0:
+        # Count detections
         class_names = [model.names[int(cls)] for cls in detections[:, 5]]
         count = Counter(class_names)
-        detection_str = ', '.join([f"{name}: {count[name]}" for name in count])
+        detection_str = '\n'.join([f"{name}: {count[name]} cells detected" for name in count])
         
-        # Extract cell attributes (position, size, etc.)
+        # Create detailed DataFrame
         df = pd.DataFrame(detections.numpy(), columns=["x_min", "y_min", "x_max", "y_max", "confidence", "class"])
         df["class_name"] = df["class"].apply(lambda x: model.names[int(x)])
         df["width"] = df["x_max"] - df["x_min"]
         df["height"] = df["y_max"] - df["y_min"]
         df["area"] = df["width"] * df["height"]
         
-        summary = df.groupby("class_name")["area"].describe().reset_index()
+        # Generate summary statistics
+        summary = df.groupby("class_name").agg({
+            'area': ['count', 'mean', 'std', 'min', 'max'],
+            'confidence': 'mean'
+        }).round(2)
+        summary.columns = ['Count', 'Mean Area', 'Std Dev', 'Min Area', 'Max Area', 'Avg Confidence']
+        summary = summary.reset_index()
+        
+        # Create visualizations
+        size_dist_plot = create_size_distribution_plot(df)
+        density_plot = create_density_heatmap(df, image.shape)
+        
+        return (
+            annotated_img,
+            detection_str,
+            summary,
+            size_dist_plot,
+            density_plot
+        )
     else:
-        detection_str = "No detections"
-        summary = pd.DataFrame(columns=["class_name", "count", "mean", "std", "min", "25%", "50%", "75%", "max"])
-    
-    return annotated_img, detection_str, summary
+        return (
+            annotated_img,
+            "No cells detected",
+            pd.DataFrame(),
+            None,
+            None
+        )
 
-# Create Gradio interface
-app = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="numpy", label="Upload an Image"),
-    outputs=[
-        gr.Image(type="numpy", label="Annotated Image"),
-        gr.Textbox(label="Detection Counts"),
-        gr.Dataframe(label="Cell Statistics")
-    ],
-    title="Bioengineering Image Analysis Tool",
-    description="Upload an image to detect and analyze bioengineering cells using YOLOv10."
-)
+# Create Gradio interface with improved layout
+with gr.Blocks(theme=gr.themes.Soft()) as app:
+    gr.Markdown("""
+    # Bioengineering Image Analysis Tool
+    Upload microscopy images to detect and analyze cells using YOLOv10.
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="numpy", label="Upload Image")
+            conf_slider = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.25,
+                step=0.05,
+                label="Confidence Threshold",
+                info="Adjust detection sensitivity"
+            )
+            analyze_btn = gr.Button("Analyze Image", variant="primary")
+        
+        with gr.Column(scale=1):
+            output_image = gr.Image(type="numpy", label="Detected Cells")
+            detection_text = gr.Textbox(label="Detection Summary", lines=3)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            stats_df = gr.Dataframe(
+                label="Cell Statistics",
+                headers=['Cell Type', 'Count', 'Mean Area', 'Std Dev', 'Min Area', 'Max Area', 'Avg Confidence']
+            )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            size_plot = gr.Plot(label="Cell Size Distribution")
+        with gr.Column(scale=1):
+            density_plot = gr.Plot(label="Cell Density Heatmap")
+    
+    # Handle button click
+    analyze_btn.click(
+        process_image,
+        inputs=[input_image, conf_slider],
+        outputs=[output_image, detection_text, stats_df, size_plot, density_plot]
+    )
+    
+    gr.Markdown("""
+    ### Instructions:
+    1. Upload a microscopy image containing cells
+    2. Adjust the confidence threshold if needed (higher values = stricter detection)
+    3. Click 'Analyze Image' to process
+    4. View results in the various panels:
+        - Annotated image shows detected cells
+        - Summary provides cell counts
+        - Statistics table shows detailed measurements
+        - Plots visualize size distribution and spatial density
+    """)
 
+# Launch the app
 app.launch()