Gradio UI added

app.py

@@ -0,0 +1,130 @@
+# app.py (in the root directory)
+
+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
+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]]
+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:
+        raise gr.Error("Patient Name and Age are required.")
+    if not image_list:
+        raise gr.Error("Please upload at least one image.")
+    
+    # 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
+    
+    return confidences
+
+# --- Function to fetch and format database records ---
+async def get_records_html():
+    records = await get_all_records()
+    if not records:
+        return "<p>No records found in the database.</p>"
+    
+    # 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(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]
+    )
+    
+    # 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)
+
+# --- Launch the App ---
+if __name__ == "__main__":
+    demo.launch()

app/database.py

@@ -0,0 +1,63 @@
+# app/database.py
+
+import os
+from motor.motor_asyncio import AsyncIOMotorClient
+from dotenv import load_dotenv
+import datetime
+from typing import List, Dict
+
+# Load environment variables from .env file
+load_dotenv()
+
+# --- Database Connection ---
+# Get the connection string from the environment variables
+MONGODB_URL = os.getenv("MONGODB_CONNECTION_STRING")
+
+if not MONGODB_URL:
+    raise ValueError("MONGODB_CONNECTION_STRING not found in environment variables. Please check your .env file.")
+
+# Create a client instance
+client = AsyncIOMotorClient(MONGODB_URL)
+
+# Get a handle to the database (it will be created if it doesn't exist)
+# The database name 'pneumonia_db' should match the one in your connection string
+database = client.pneumonia_db
+
+# Get a handle to the collection (like a table in SQL)
+patient_collection = database.get_collection("patient_records")
+
+
+# --- Database Operations (now async) ---
+
+async def add_patient_record(name: str, age: int, result: str, confidence: float) -> Dict:
+    """
+    Inserts a new patient record into the MongoDB collection.
+    
+    Returns the inserted document.
+    """
+    record_document = {
+        "name": name,
+        "age": age,
+        "prediction_result": result,
+        "confidence_score": confidence,
+        "timestamp": datetime.datetime.utcnow()
+    }
+    
+    # .insert_one is an async operation, so we must 'await' it
+    result = await patient_collection.insert_one(record_document)
+    
+    # Find the newly created document to return it
+    new_record = await patient_collection.find_one({"_id": result.inserted_id})
+    return new_record
+
+
+async def get_all_records() -> List[Dict]:
+    """
+    Retrieves all patient records, sorted by the most recent timestamp.
+    """
+    records = []
+    # .find() returns a cursor, which we iterate over asynchronously
+    cursor = patient_collection.find({}).sort("timestamp", -1) # -1 for descending order
+    async for document in cursor:
+        records.append(document)
+    return records

app/prediction.py

@@ -0,0 +1,53 @@
+# app/prediction.py
+
+import torch
+from transformers import ViTImageProcessor, ViTForImageClassification
+from PIL import Image
+from pathlib import Path
+import numpy as np
+from typing import List, Dict, Union
+
+# Define a type hint for the input, which can be a path or bytes
+ImageType = Union[str, Path, bytes]
+
+class PredictionPipeline:
+    def __init__(self, model_path: Path = Path("artifacts/model_training/model")):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.processor = ViTImageProcessor.from_pretrained(model_path)
+        self.model = ViTForImageClassification.from_pretrained(model_path).to(self.device)
+        self.model.eval()
+        self.id2label = self.model.config.id2label
+
+    def predict(self, image_sources: List[ImageType]) -> Dict[str, Union[str, float]]:
+        if not image_sources:
+            return {"prediction": "Error", "confidence": 0.0, "details": "No images provided."}
+
+        all_logits = []
+        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")
+                
+                inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+                
+                with torch.no_grad():
+                    outputs = self.model(**inputs)
+                    all_logits.append(outputs.logits)
+            except Exception as e:
+                print(f"Skipping a corrupted or invalid image file. Error: {e}")
+                continue
+        
+        if not all_logits:
+             return {"prediction": "Error", "confidence": 0.0, "details": "All provided images were invalid."}
+
+        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()]
+        
+        return {
+            "prediction": predicted_label,
+            "confidence": confidence_score.item()
+        }

requirements.txt

@@ -18,4 +18,10 @@ dvc
 matplotlib
 Pillow
 kaggle
-python-dotenv
+python-dotenv
+nicegui
+sqlalchemy
+pymongo
+motor
+huggingface_hub
+gradio

src/vitClassifier/pipeline/prediction.py

@@ -0,0 +1,100 @@
+# prediction.py
+
+import torch
+from transformers import ViTImageProcessor, ViTForImageClassification
+from PIL import Image
+import argparse
+import os
+from pathlib import Path
+
+class PredictionPipeline:
+    def __init__(self, model_path: str = "artifacts/model_training/model"):
+        """
+        Initializes the prediction pipeline by loading the trained model and processor.
+        
+        Args:
+            model_path (str): The path to the directory containing the saved model and processor.
+        """
+        # Set the device
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+        
+        # Load the processor and model from the specified path
+        self.processor = ViTImageProcessor.from_pretrained(model_path)
+        self.model = ViTForImageClassification.from_pretrained(model_path).to(self.device)
+        self.model.eval() # Set the model to evaluation mode
+        
+        # Get the label mappings from the model's configuration
+        self.id2label = self.model.config.id2label
+
+    def predict(self, image_path: str):
+        """
+        Makes a prediction on a single image.
+        
+        Args:
+            image_path (str): The file path of the image to be classified.
+            
+        Returns:
+            dict: A dictionary containing the predicted label and its confidence score.
+        """
+        try:
+            # Open the image using PIL (Python Imaging Library)
+            image = Image.open(image_path).convert("RGB")
+        except FileNotFoundError:
+            return {"error": f"Image not found at path: {image_path}"}
+        except Exception as e:
+            return {"error": f"Failed to open image: {e}"}
+
+        # Preprocess the image using the ViTImageProcessor
+        # This handles resizing, normalization, and conversion to a tensor
+        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+        
+        # Make a prediction
+        with torch.no_grad(): # Disable gradient calculation for faster inference
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+
+        # Get the predicted class index
+        predicted_class_idx = logits.argmax(-1).item()
+        
+        # Get the human-readable label
+        predicted_label = self.id2label[predicted_class_idx]
+        
+        # Calculate the confidence score using softmax
+        probabilities = torch.nn.functional.softmax(logits, dim=-1)
+        confidence_score = probabilities[0][predicted_class_idx].item()
+        
+        result = {
+            "predicted_label": predicted_label,
+            "confidence_score": f"{confidence_score:.4f}"
+        }
+        
+        return result
+
+if __name__ == '__main__':
+    # --- How to run this script from the command line ---
+    # Example 1 (Pneumonia):
+    # python prediction.py --image "artifacts/data_ingestion/chest_xray/test/PNEUMONIA/person100_bacteria_475.jpeg"
+    
+    # Example 2 (Normal):
+    # python prediction.py --image "artifacts/data_ingestion/chest_xray/test/NORMAL/IM-0001-0001.jpeg"
+
+    # Set up argument parser to accept image path from the command line
+    parser = argparse.ArgumentParser(description="Chest X-ray Pneumonia Detection")
+    parser.add_argument("--image", type=str, required=True, help="Path to the input image")
+    args = parser.parse_args()
+
+    # Create an instance of the pipeline
+    pipeline = PredictionPipeline()
+    
+    # Make a prediction
+    result = pipeline.predict(args.image)
+    
+    # Print the result
+    print("\n--- Prediction Result ---")
+    if "error" in result:
+        print(f"Error: {result['error']}")
+    else:
+        print(f"The model predicts this is a '{result['predicted_label']}' case.")
+        print(f"Confidence: {result['confidence_score']}")
+    print("-------------------------\n")