Prepare application for deployment

README.md

@@ -1 +1,35 @@
-# End-to-End-Chest-X-ray-Pneumonia-Detection-with-ViT
+---
+title: Pneumonia Detection AI
+emoji: 🩺
+colorFrom: blue
+colorTo: indigo
+sdk: gradio
+sdk_version: 4.19.1
+app_file: app.py
+pinned: false
+---
+
+# 🩺 Pneumonia Detection AI
+
+This Space demonstrates a complete, end-to-end MLOps pipeline for medical image classification.
+
+## ✨ Features
+
+-   **AI-Powered Diagnosis:** Upload one or more chest X-ray images to get an instant classification of **Normal** or **Pneumonia**.
+-   **Advanced Model:** Powered by a fine-tuned **Vision Transformer (ViT)** for high accuracy.
+-   **Multi-Image Analysis:** The AI provides both an overall prediction for the patient and individual watermarked results for each image.
+-   **Patient History:** All analyses are logged to a **MongoDB** database and can be reviewed.
+-   **Sample Library:** Test the app instantly with a library of sample X-ray images.
+
+## 🛠️ Tech Stack
+
+-   **Model:** Google's `vit-base-patch16-224-in21k`
+-   **MLOps Pipeline:** DVC & MLflow
+-   **Frontend:** Gradio
+-   **Database:** MongoDB Atlas
+-   **Hosting:** Hugging Face Spaces
+
+This project was developed by **Alyyan Ahmed** and **Munim Akbar**.
+
+---
+**Disclaimer:** This is a demo application for educational and portfolio purposes. It is **not a certified medical device** and should not be used for actual medical diagnosis.

app.py

@@ -1,129 +1,181 @@
-# app.py (in the root directory)
+# app.py (Final UI Polish Version)
 
 import gradio as gr
 from pathlib import Path
 from huggingface_hub import snapshot_download
 import asyncio
-from PIL import Image
 
-# --- Import and Initialize Backend Components from the 'app' folder ---
 from app.prediction import PredictionPipeline
 from app.database import add_patient_record, get_all_records
 
-# Initialize components once
+# --- Initialization ---
 prediction_pipeline = PredictionPipeline()
 HF_DATASET_REPO = "ALYYAN/chest-xray-pneumonia-samples"
 try:
     SAMPLE_IMAGE_DIR = Path(snapshot_download(repo_id=HF_DATASET_REPO, repo_type="dataset"))
-    # Create a list of sample image paths for the Gradio component
-    SAMPLE_IMAGES = [str(p) for p in list(SAMPLE_IMAGE_DIR.glob('*/*.jpeg'))[:10]]
+    SAMPLE_IMAGES = [str(p) for p in list(SAMPLE_IMAGE_DIR.glob('*/*.jpeg'))]
 except Exception as e:
     print(f"Could not download sample images: {e}")
     SAMPLE_IMAGES = []
 
-# --- Core Prediction Logic for Gradio ---
-async def classify_images(patient_name, patient_age, image_list):
-    # 1. Input Validation
-    if not patient_name or patient_age is None:
+# --- Core Logic (Async Functions) ---
+async def process_analysis(patient_name, patient_age, image_list, is_sample=False):
+    if not is_sample and (not patient_name or patient_age is None or str(patient_age).strip() == ""):
         raise gr.Error("Patient Name and Age are required.")
     if not image_list:
-        raise gr.Error("Please upload at least one image.")
+        raise gr.Error("At least one image is required.")
     
-    # Gradio provides file paths for uploaded files in a temp directory
-    # Our prediction pipeline can handle these paths directly.
-
-    # 2. Run Prediction
-    result = prediction_pipeline.predict(image_list) # Pass the list of temp file paths
-    prediction = result.get("prediction", "Error")
-    confidence = result.get("confidence", 0)
-
-    if prediction == "Error":
-        raise gr.Error(result.get("details", "An unknown error occurred during prediction."))
-
-    # 3. Save to Database
-    # Ensure age is an integer
-    try:
-        age = int(patient_age)
-    except (ValueError, TypeError):
-        raise gr.Error("Patient Age must be a valid number.")
-
-    await add_patient_record(
-        name=str(patient_name),
-        age=age,
-        result=prediction,
-        confidence=confidence
-    )
-
-    # 4. Format the Output for Gradio
-    confidences = {"NORMAL": 0.0, "PNEUMONIA": 0.0} # Initialize both labels
-    confidences[prediction] = confidence
+    result = prediction_pipeline.predict(image_list)
+    if "error" in result:
+        raise gr.Error(result["error"])
+
+    final_pred = result["final_prediction"]
+    final_conf = result["final_confidence"]
     
-    return confidences
+    if not is_sample:
+        await add_patient_record(str(patient_name), int(patient_age), final_pred, final_conf)
 
-# --- Function to fetch and format database records ---
-async def get_records_html():
+    confidences = {"NORMAL": 0.0, "PNEUMONIA": 0.0}
+    confidences[final_pred] = final_conf
+    confidences["NORMAL" if final_pred == "PNEUMONIA" else "PNEUMONIA"] = 1 - final_conf
+    
+    return [
+        gr.update(visible=False), # uploader_column
+        gr.update(visible=True),  # results_column
+        gr.update(value=result["watermarked_images"]), # result_images
+        gr.update(value=confidences) # result_label
+    ]
+
+async def refresh_history_table():
     records = await get_all_records()
-    if not records:
-        return "<p>No records found in the database.</p>"
+    data_for_df = []
+    if records:
+        data_for_df = [[r.get('name'), r.get('age'), r.get('prediction_result'), f"{r.get('confidence_score', 0):.2%}", r.get('timestamp').strftime('%Y-%m-%d %H:%M')] for r in records]
+    return gr.update(value=data_for_df)
+
+# --- Gradio UI Definition ---
+css = """
+/* --- Professional Dark Theme & Fonts --- */
+:root { --primary-hue: 220 !important; --secondary-hue: 210 !important; --neutral-hue: 210 !important; --body-background-fill: #111827 !important; --block-background-fill: #1F2937 !important; --block-border-width: 1px !important; --border-color-accent: #374151 !important; --background-fill-secondary: #1F2937 !important;}
+/* --- Header & Title Styling --- */
+#app_header { text-align: center; }
+#app_title { font-size: 2.8rem !important; font-weight: 700 !important; color: #FFFFFF !important; padding-top: 1rem; }
+#app_subtitle { font-size: 1.2rem !important; color: #9CA3AF !important; margin-bottom: 2rem; }
+/* --- Layout, Spacing, and Component Styling --- */
+#main_container { gap: 2rem; }
+#results_gallery { height: 350px !important; }
+#results_gallery .gallery-item { height: 330px !important; max-height: 330px !important; padding: 0.25rem !important; background-color: #374151; border: 1px solid #374151 !important; }
+#results_gallery .gallery-item img { object-fit: contain !important; }
+#bottom_controls { max-width: 600px; margin: 2.5rem auto 1rem auto; }
+#bottom_controls .gr-accordion > .gr-block-label { text-align: center !important; display: block !important; }
+"""
+with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="blue"), css=css, title="Pneumonia Detection AI") as demo:
     
-    # Create an HTML table from the records
-    html = "<table><tr><th>Name</th><th>Age</th><th>Prediction</th><th>Confidence</th><th>Date</th></tr>"
-    for r in records:
-        confidence_percent = f"{r['confidence_score']:.2%}" if r['confidence_score'] is not None else "N/A"
-        timestamp = r['timestamp'].strftime('%Y-%m-%d %H:%M') if r['timestamp'] else "N/A"
-        html += f"<tr><td>{r.get('name', 'N/A')}</td><td>{r.get('age', 'N/A')}</td><td>{r.get('prediction_result', 'N/A')}</td><td>{confidence_percent}</td><td>{timestamp}</td></tr>"
-    html += "</table>"
-    return html
-
-# --- Build the Gradio Interface ---
-with gr.Blocks(theme=gr.themes.Soft(), css="table { width: 100%; border-collapse: collapse; } th, td { padding: 8px; text-align: left; border-bottom: 1px solid #ddd; }") as demo:
-    gr.Markdown("# 🩺 Pneumonia Detection AI")
-    gr.Markdown("Upload one or more chest X-ray images for a patient to classify them as **Normal** or **Pneumonia**.")
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### 1. Patient Information")
-            patient_name = gr.Textbox(label="Patient Name", placeholder="e.g., John Doe")
-            patient_age = gr.Number(label="Patient Age", minimum=0, maximum=120, step=1)
-
-            gr.Markdown("### 2. Upload Images")
-            # Using type="filepath" is simpler and avoids memory issues with large images
-            image_input = gr.File(
-                label="Upload up to 3 X-Rays",
-                file_count="multiple",
-                file_types=["image"],
-                type="filepath" # Gradio will save uploads to a temp dir and give us the path
-            )
-            
-            submit_btn = gr.Button("Analyze Images", variant="primary")
-
-            if SAMPLE_IMAGES:
-                gr.Examples(
-                    examples=SAMPLE_IMAGES,
-                    inputs=image_input,
-                    label="Sample Images (Click one, then click Analyze)",
-                    examples_per_page=5
+    with gr.Column() as main_app:
+        with gr.Column(elem_id="app_header"):
+            gr.Markdown("# 🩺 Pneumonia Detection AI", elem_id="app_title")
+            gr.Markdown("An AI-powered tool to assist in the diagnosis of pneumonia.", elem_id="app_subtitle")
+        with gr.Row(elem_id="main_container"):
+            with gr.Column(scale=1) as uploader_column:
+                gr.Markdown("### Upload Patient X-Rays")
+                image_input = gr.File(label="Upload up to 3 Images", file_count="multiple", file_types=["image"], type="filepath")
+            with gr.Column(scale=2, visible=False) as results_column:
+                gr.Markdown("### Analysis Results")
+                result_images = gr.Gallery(label="Analyzed Images", columns=3, object_fit="contain", height=350, elem_id="results_gallery")
+                result_label = gr.Label(label="Overall Prediction", num_top_classes=2)
+                start_over_btn = gr.Button("Start New Analysis", variant="secondary")
+        with gr.Group(visible=False) as patient_info_modal:
+            gr.Markdown("## Enter Patient Details", elem_classes="text-center")
+            patient_name_modal = gr.Textbox(label="Patient Name", placeholder="e.g., John Doe")
+            patient_age_modal = gr.Number(label="Patient Age", minimum=0, maximum=120, step=1)
+            with gr.Row():
+                submit_analysis_btn = gr.Button("Analyze Images", variant="primary")
+                cancel_btn = gr.Button("Cancel", variant="stop")
+        with gr.Column(elem_id="bottom_controls"):
+            with gr.Accordion("About this Tool", open=False):
+                gr.Markdown(
+                    """
+                    ### MLOps-Powered Pneumonia Detection
+
+                    This application demonstrates a complete, end-to-end MLOps pipeline for medical image classification. It leverages a state-of-the-art **Vision Transformer (ViT)** model, fine-tuned on a public dataset of chest X-ray images to distinguish between Normal and Pneumonia cases.
+
+                    ---
+
+                    **Key Features & Technologies:**
+
+                    *   **Model:** Google's `vit-base-patch16-224-in21k`, fine-tuned for high accuracy.
+                    *   **MLOps Pipeline:** Reproducible workflow managed by **DVC** for data versioning and **MLflow** for experiment tracking.
+                    *   **Database:** Patient and prediction data is stored and managed in a **MongoDB** database for scalability.
+                    *   **Frontend:** A responsive and interactive user interface built with **Gradio**.
+                    *   **Deployment Ready:** The entire project is containerized and ready for deployment on platforms like Hugging Face Spaces.
+
+                    **Disclaimer:** This tool is for demonstration and educational purposes only and is **not a substitute for professional medical advice.**
+
+                    ---
+
+                    **Project Team:**
+
+                    *   **Alyyan Ahmed** - (roles)
+                    *   **Munim Akbar** - (roles)
+                    """
                 )
-
-        with gr.Column(scale=1):
-            gr.Markdown("### 3. Analysis Results")
-            output_label = gr.Label(label="Prediction", num_top_classes=2)
-            gr.Markdown("---")
-            with gr.Accordion("View Patient Record History", open=False):
-                records_html = gr.HTML("Loading records...")
-                demo.load(get_records_html, None, records_html) # Load records when the app starts
-                refresh_btn = gr.Button("Refresh History")
-
-
-    # --- Link Components to the Function ---
-    submit_btn.click(
-        fn=classify_images,
-        inputs=[patient_name, patient_age, image_input],
-        outputs=[output_label]
-    )
+            with gr.Row():
+                samples_btn = gr.Button("Try Sample Images")
+                history_btn = gr.Button("View Patient History")
+    with gr.Column(visible=False) as history_page:
+        gr.Markdown("# 📜 Patient Record History", elem_classes="app_title")
+        with gr.Row():
+            back_to_main_btn_hist = gr.Button("⬅️ Back to Main App")
+            refresh_history_btn = gr.Button("Refresh History")
+        history_df = gr.DataFrame(headers=["Name", "Age", "Prediction", "Confidence", "Date"], row_count=10, interactive=False)
+    with gr.Column(visible=False) as samples_page:
+        gr.Markdown("# 🖼️ Sample Image Library", elem_classes="app_title")
+        gr.Markdown("Click an image to run an anonymous analysis.")
+        back_to_main_btn_samp = gr.Button("⬅️ Back to Main App")
+        sample_gallery = gr.Gallery(value=SAMPLE_IMAGES, label="Sample Images", columns=5, height=400)
+    
+    # --- Event Handling Logic ---
+    def show_patient_info(files):
+        return gr.update(visible=True) if files else gr.update(visible=False)
+    image_input.upload(fn=show_patient_info, inputs=image_input, outputs=patient_info_modal)
+    
+    async def submit_and_hide_modal(name, age, files):
+        analysis_results = await process_analysis(name, age, files)
+        return [
+            *analysis_results,
+            gr.update(visible=False) # Hide the modal
+        ]
+    submit_analysis_btn.click(fn=submit_and_hide_modal, inputs=[patient_name_modal, patient_age_modal, image_input], outputs=[uploader_column, results_column, result_images, result_label, patient_info_modal])
+    
+    cancel_btn.click(lambda: (gr.update(visible=False), None), None, [patient_info_modal, image_input])
+    start_over_btn.click(fn=None, js="() => { window.location.reload(); }")
+    
+    async def handle_sample_click(evt: gr.SelectData):
+        selected_path = evt.value
+        analysis_results = await process_analysis("Sample User", 0, [selected_path], is_sample=True)
+        return [
+            gr.update(visible=True),   # main_app
+            gr.update(visible=False),  # samples_page
+            *analysis_results          
+        ]
+    sample_gallery.select(handle_sample_click, None, [main_app, samples_page, uploader_column, results_column, result_images, result_label])
+    
+    all_pages = [main_app, history_page, samples_page]
+    async def show_history_page_and_refresh():
+        records_update = await refresh_history_table()
+        return [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), records_update]
+    def show_samples_page():
+        return [gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)]
+    def show_main_page():
+        return [gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)]
+    
+    history_btn.click(fn=show_history_page_and_refresh, outputs=all_pages + [history_df])
+    samples_btn.click(fn=show_samples_page, outputs=all_pages)
+    back_to_main_btn_hist.click(fn=show_main_page, outputs=all_pages)
+    back_to_main_btn_samp.click(fn=show_main_page, outputs=all_pages)
     
-    # When the refresh button is clicked, re-run the get_records_html function
-    refresh_btn.click(fn=get_records_html, inputs=None, outputs=records_html)
+    refresh_history_btn.click(fn=refresh_history_table, outputs=history_df)
+    demo.load(fn=refresh_history_table, outputs=history_df)
 
 # --- Launch the App ---
 if __name__ == "__main__":

app/image_utils.py

@@ -0,0 +1,65 @@
+# app/image_utils.py
+
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+def add_watermark(image_array: np.ndarray, text: str, confidence: float) -> Image.Image:
+    """
+    Adds a translucent watermark to an image with the prediction result and confidence.
+
+    Args:
+        image_array: The input image as a NumPy array.
+        text: The prediction text (e.g., "NORMAL" or "PNEUMONIA").
+        confidence: The confidence score of the prediction.
+
+    Returns:
+        A PIL Image object with the watermark applied.
+    """
+    # Convert NumPy array to PIL Image
+    image = Image.fromarray(image_array).convert("RGBA")
+    
+    # Create a transparent overlay for the text
+    txt_overlay = Image.new("RGBA", image.size, (255, 255, 255, 0))
+    draw = ImageDraw.Draw(txt_overlay)
+
+    # Define watermark properties
+    is_pneumonia = (text == "PNEUMONIA")
+    box_color = (220, 53, 69, 180) if is_pneumonia else (25, 135, 84, 180) # Red for Pneumonia, Green for Normal
+    text_color = (255, 255, 255, 255)
+    
+    # Define font (uses a default if a specific .ttf is not found)
+    try:
+        font_size = int(image.height / 8)
+        font = ImageFont.truetype("arialbd.ttf", font_size)
+    except IOError:
+        print("Arial Bold font not found, using default. Watermark quality may be lower.")
+        font_size = int(image.height / 8)
+        font = ImageFont.load_default()
+
+    # Define text and box position
+    text_to_draw = f"{text}\n{confidence:.1%}"
+    
+    # Get text size
+    try:
+        # Use getbbox for modern Pillow versions
+        _, _, text_width, text_height = draw.textbbox((0, 0), text_to_draw, font=font)
+    except AttributeError:
+        # Fallback for older Pillow versions
+        text_width, text_height = draw.textsize(text_to_draw, font=font)
+        
+    position = (20, 20) # Top-left corner with some padding
+    box_position = [
+        position[0] - 10,
+        position[1] - 10,
+        position[0] + text_width + 10,
+        position[1] + text_height + 10
+    ]
+
+    # Draw the semi-transparent rectangle and the text
+    draw.rectangle(box_position, fill=box_color)
+    draw.text(position, text_to_draw, font=font, fill=text_color)
+    
+    # Combine the overlay with the original image
+    watermarked_image = Image.alpha_composite(image, txt_overlay)
+    
+    return watermarked_image.convert("RGB")

app/prediction.py

@@ -5,10 +5,10 @@ from transformers import ViTImageProcessor, ViTForImageClassification
 from PIL import Image
 from pathlib import Path
 import numpy as np
-from typing import List, Dict, Union
+from typing import List, Dict, Union, Any
+from .image_utils import add_watermark
 
-# Define a type hint for the input, which can be a path or bytes
-ImageType = Union[str, Path, bytes]
+ImageType = Union[str, Path, bytes, np.ndarray]
 
 class PredictionPipeline:
     def __init__(self, model_path: Path = Path("artifacts/model_training/model")):
@@ -18,36 +18,63 @@ class PredictionPipeline:
         self.model.eval()
         self.id2label = self.model.config.id2label
 
-    def predict(self, image_sources: List[ImageType]) -> Dict[str, Union[str, float]]:
+    def predict(self, image_sources: List[ImageType]) -> Dict[str, Any]:
         if not image_sources:
-            return {"prediction": "Error", "confidence": 0.0, "details": "No images provided."}
+            return {"error": "No images provided."}
 
+        individual_results = []
         all_logits = []
+        valid_images_as_np = []
+
         for source in image_sources:
             try:
-                # --- THIS IS THE FIX ---
-                # The Image.open() function can handle both paths and byte streams.
-                # No special handling is needed.
-                image = Image.open(source).convert("RGB")
+                if isinstance(source, np.ndarray):
+                    image = Image.fromarray(source).convert("RGB")
+                else:
+                    image = Image.open(source).convert("RGB")
                 
-                inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+                valid_images_as_np.append(np.array(image))
                 
+                inputs = self.processor(images=image, return_tensors="pt").to(self.device)
                 with torch.no_grad():
                     outputs = self.model(**inputs)
-                    all_logits.append(outputs.logits)
+                    logits = outputs.logits
+                    all_logits.append(logits)
+                    
+                    # --- NEW: Calculate individual prediction ---
+                    ind_probs = torch.nn.functional.softmax(logits, dim=-1)
+                    ind_conf, ind_idx = torch.max(ind_probs, dim=-1)
+                    individual_results.append({
+                        "prediction": self.id2label[ind_idx.item()],
+                        "confidence": ind_conf.item()
+                    })
+
             except Exception as e:
                 print(f"Skipping a corrupted or invalid image file. Error: {e}")
+                individual_results.append({"prediction": "Error", "confidence": 0})
                 continue
         
         if not all_logits:
-             return {"prediction": "Error", "confidence": 0.0, "details": "All provided images were invalid."}
+             return {"error": "All images were invalid."}
 
+        # --- Aggregate Prediction (same as before) ---
         avg_logits = torch.mean(torch.stack(all_logits), dim=0)
         probabilities = torch.nn.functional.softmax(avg_logits, dim=-1)
         confidence_score, predicted_class_idx = torch.max(probabilities, dim=-1)
-        predicted_label = self.id2label[predicted_class_idx.item()]
+        
+        final_prediction = self.id2label[predicted_class_idx.item()]
+        final_confidence = confidence_score.item()
+
+        # --- NEW: Watermark images with their INDIVIDUAL results ---
+        watermarked_images = [
+            add_watermark(img_np, res["prediction"], res["confidence"])
+            for img_np, res in zip(valid_images_as_np, individual_results)
+            if res["prediction"] != "Error"
+        ]
         
         return {
-            "prediction": predicted_label,
-            "confidence": confidence_score.item()
+            "final_prediction": final_prediction,
+            "final_confidence": final_confidence,
+            "individual_results": individual_results,
+            "watermarked_images": watermarked_images
         }

requirements.txt

@@ -1,27 +1,16 @@
-pandas
-numpy
-torch
-torchvision
+gradio==4.19.1
+pymongo
+motor
+python-dotenv
+huggingface_hub
+torch --index-url https://download.pytorch.org/whl/cpu
+torchvision --index-url https://download.pytorch.org/whl/cpu
+Pillow
 transformers
-datasets>=2.14.5
-evaluate
-accelerate>=0.27
-mlflow
+datasets
 scikit-learn
 imblearn
 python-box
 PyYAML
 ensure
-tqdm
-pathlib
-dvc
-matplotlib
-Pillow
-kaggle
-python-dotenv
-nicegui
-sqlalchemy
-pymongo
-motor
-huggingface_hub
-gradio
+dvc[gdrive] # Add dvc with gdrive support