Update app.py

app.py

@@ -7,12 +7,16 @@ from datasets import load_dataset, DownloadConfig
 import matplotlib.pyplot as plt
 import numpy as np
 import cv2
+from groq import Groq
+
+# Initialize Groq client
+client = Groq(api_key="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")
 
 # Model and processor configuration
 model_name_or_path = "google/vit-base-patch16-224-in21k"
 processor = ViTImageProcessor.from_pretrained(model_name_or_path)
 
-# Load dataset (adjust dataset_path accordingly)
+# Load dataset
 dataset_path = "pawlo2013/chest_xray"
 download_config = DownloadConfig(max_retries=10)
 train_dataset = load_dataset(dataset_path, split="train", download_config=download_config)
@@ -26,126 +30,139 @@ model = ViTForImageClassification.from_pretrained(
     label2id={label: i for i, label in enumerate(class_names)},
 )
 
-# Set model to evaluation mode
 model.eval()
 
-
-# Define the classification function
-def classify_and_visualize(img, device="cpu", discard_ratio=0.9, head_fusion="mean"):
-    img = img.convert("RGB")
-    processed_input = processor(images=img, return_tensors="pt").to(device)
-    processed_input = processed_input["pixel_values"].to(device)
-
-    with torch.no_grad():
-        outputs = model(processed_input, output_attentions=True)
-        logits = outputs.logits
-        probabilities = torch.softmax(logits, dim=1)[0].tolist()
-        prediction = torch.argmax(logits, dim=-1).item()
-        predicted_class = class_names[prediction]
-
-    result = {class_name: prob for class_name, prob in zip(class_names, probabilities)}
-
-    # Generate attention heatmap
-    heatmap_img = show_final_layer_attention_maps(
-        outputs, processed_input, device, discard_ratio, head_fusion
+def get_ai_explanation(diagnosis, probabilities):
+    if diagnosis == "normal":
+        prompt = f"""Given a chest X-ray analysis showing NORMAL results with {probabilities['normal']:.2%} confidence:
+        1. Explain what this means
+        2. Suggest when they should still consider consulting a doctor
+        3. List key symptoms that would warrant medical attention
+        Keep the tone professional yet reassuring."""
+    else:
+        prompt = f"""Given a chest X-ray analysis showing {diagnosis} pneumonia with {probabilities[diagnosis]:.2%} confidence:
+        1. Explain what {diagnosis} pneumonia is
+        2. List immediate steps the patient should take
+        3. Provide care recommendations
+        4. Mention warning signs to watch for
+        Keep the tone informative and caring but emphasize the importance of professional medical consultation."""
+
+    completion = client.chat.completions.create(
+        messages=[{"role": "user", "content": prompt}],
+        model="mixtral-8x7b-32768",
+        temperature=0.7,
     )
-
-    return {"probabilities": result, "heatmap": heatmap_img}
-
-
-def format_output(output):
-    return (output["probabilities"], output["heatmap"])
-
-
-# Function to load examples from a folder
-def load_examples_from_folder(folder_path):
-    examples = []
-    for file in os.listdir(folder_path):
-        if file.endswith((".png", ".jpg", ".jpeg")):
-            examples.append(Image.open(os.path.join(folder_path, file)))
-    return examples
-
-
-# Function to show final layer attention maps
-def show_final_layer_attention_maps(
-    outputs,
-    processed_input,
-    device,
-    discard_ratio=0.6,
-    head_fusion="max",
-    only_last_layer=False,
-):
-
-    with torch.no_grad():
-        image = processed_input.squeeze(0)
-        image = image - image.min()
-        image = image / image.max()
-
-        result = torch.eye(outputs.attentions[0].size(-1)).to(device)
-        if only_last_layer:
-            attention_list = outputs.attentions[-1].unsqueeze(0).to(device)
-        else:
-            attention_list = outputs.attentions
-
-        for attention in attention_list:
-            if head_fusion == "mean":
-                attention_heads_fused = attention.mean(axis=1)
-            elif head_fusion == "max":
-                attention_heads_fused = attention.max(axis=1)[0]
-            elif head_fusion == "min":
-                attention_heads_fused = attention.min(axis=1)[0]
-
-            flat = attention_heads_fused.view(attention_heads_fused.size(0), -1)
-            _, indices = flat.topk(int(flat.size(-1) * discard_ratio), -1, False)
-            indices = indices[indices != 0]
-            flat[0, indices] = 0
-
-            I = torch.eye(attention_heads_fused.size(-1)).to(device)
-            a = (attention_heads_fused + 1.0 * I) / 2
-            a = a / a.sum(dim=-1)
-            result = torch.matmul(a, result)
-
-        mask = result[0, 0, 1:]
-        width = int(mask.size(-1) ** 0.5)
-        mask = mask.reshape(width, width).cpu().numpy()
-        mask = mask / np.max(mask)
-
-        mask = cv2.resize(mask, (224, 224))
-        mask = (mask - np.min(mask)) / (np.max(mask) - np.min(mask))
-        heatmap = plt.cm.jet(mask)[:, :, :3]
-
-        showed_img = image.permute(1, 2, 0).detach().cpu().numpy()
-        showed_img = (showed_img - np.min(showed_img)) / (
-            np.max(showed_img) - np.min(showed_img)
-        )
-        superimposed_img = heatmap * 0.4 + showed_img * 0.6
-
-        superimposed_img_pil = Image.fromarray(
-            (superimposed_img * 255).astype(np.uint8)
-        )
-        return superimposed_img_pil
-
-
-# Define the path to the examples folder
-examples_folder = "./Examples"
-examples = load_examples_from_folder(examples_folder)
-
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=lambda img: format_output(classify_and_visualize(img)),
-    inputs=gr.Image(type="pil", label="Upload X-Ray Image"),
-    outputs=[
-        gr.Label(),
-        gr.Image(label="Attention Heatmap"),
-    ],
-    examples=examples,
-    cache_examples=False,
-    flagging_mode="never",
-    concurrency_limit=1,
-    title="Pneumonia X-Ray 3-Class Classification with Vision Transformer (ViT) using data augmentation",
-    description="Upload an X-ray image to classify it as normal, viral or bacterial pneumonia. Checkout the model in more details [here](https://huggingface.co/pawlo2013/vit-pneumonia-x-ray_3_class). The examples presented are taken from the test set of [Kermany et al. (2018) dataset.](https://data.mendeley.com/datasets/rscbjbr9sj/2.) The attention heatmap over all layers of the transfomer done by the attention rollout techinique by the implementation of [jacobgil](https://github.com/jacobgil/vit-explain).",
-)
+    
+    return completion.choices[0].message.content
+
+# Rest of your existing functions (classify_and_visualize, show_final_layer_attention_maps, etc.) remain the same
+[Previous functions remain unchanged...]
+
+def create_interface():
+    # Custom CSS
+    custom_css = """
+    .logo-container { text-align: center; margin-bottom: 20px; }
+    .logo-container img { max-width: 300px; }
+    .welcome-message { text-align: center; margin: 20px 0; padding: 20px; background-color: #f5f5f5; border-radius: 10px; }
+    .model-explanation { margin: 20px 0; padding: 20px; background-color: #f0f7ff; border-radius: 10px; }
+    .pneumonia-info { margin: 20px 0; padding: 20px; background-color: #fff5f5; border-radius: 10px; }
+    .disclaimer { margin-top: 20px; padding: 20px; background-color: #f5f5f5; border-radius: 10px; font-size: 0.9em; }
+    """
+
+    # HTML Components
+    logo_html = """
+    <div class="logo-container">
+        <img src="file/logo.png" alt="PneumoInsight Logo">
+    </div>
+    """
+
+    welcome_message = """
+    <div class="welcome-message">
+        <h1>Welcome to PneumoInsight</h1>
+        <p>PneumoInsight is a side project of EarlyMed—an initiative by our team at VIT-AP University dedicated to empowering you with early health insights. 
+        Leveraging AI for early detection, our mission is simple: "Early Detection, Smarter Decision." 
+        This project is one of our key efforts to help you stay informed before visiting a doctor.</p>
+    </div>
+    """
+
+    model_explanation = """
+    <div class="model-explanation">
+        <h2>How Our Model Works</h2>
+        <p>Our system uses a Vision Transformer (ViT) model to analyze chest X-ray images. The attention heatmap visualizes 
+        areas the AI focuses on while making its diagnosis, helping make the decision-making process more transparent. 
+        The warmer colors (red/yellow) indicate areas of higher attention.</p>
+        <p>Credits: The attention heatmap visualization is implemented using the attention rollout technique by 
+        <a href="https://github.com/jacobgil/vit-explain" target="_blank">jacobgil</a>.</p>
+    </div>
+    """
+
+    pneumonia_info = """
+    <div class="pneumonia-info">
+        <h2>Understanding Pneumonia</h2>
+        <p>Pneumonia is an infection that inflames the air sacs in one or both lungs. Common symptoms include:</p>
+        <ul>
+            <li>Chest pain when breathing or coughing</li>
+            <li>Cough with phlegm or pus</li>
+            <li>Fatigue and difficulty breathing</li>
+            <li>Fever, sweating, and shaking chills</li>
+        </ul>
+        <p>Prevention tips:</p>
+        <ul>
+            <li>Get vaccinated</li>
+            <li>Practice good hygiene</li>
+            <li>Don't smoke</li>
+            <li>Maintain a strong immune system</li>
+        </ul>
+    </div>
+    """
+
+    disclaimer = """
+    <div class="disclaimer">
+        <h3>Disclaimer</h3>
+        <p>This tool is for educational purposes only and should not be used as a substitute for professional medical advice, 
+        diagnosis, or treatment. Always seek the advice of your physician or other qualified health provider.</p>
+        <p>Created by the team at VIT-AP University. View the source code on 
+        <a href="https://github.com/Mahatir-Ahmed-Tusher/PneumoInsight" target="_blank">GitHub</a>.</p>
+    </div>
+    """
+
+    def enhanced_classification(img):
+        if img is None:
+            return None, None, "Please upload an image to proceed."
+        
+        result = classify_and_visualize(img)
+        probabilities = result["probabilities"]
+        heatmap = result["heatmap"]
+        
+        # Get the predicted class
+        predicted_class = max(probabilities.items(), key=lambda x: x[1])[0]
+        
+        # Get AI explanation
+        ai_explanation = get_ai_explanation(predicted_class, probabilities)
+        
+        return probabilities, heatmap, ai_explanation
+
+    # Create the Gradio interface
+    iface = gr.Interface(
+        fn=enhanced_classification,
+        inputs=gr.Image(type="pil", label="Upload Chest X-Ray Image"),
+        outputs=[
+            gr.Label(label="Diagnosis Probabilities"),
+            gr.Image(label="Attention Heatmap"),
+            gr.Textbox(label="AI Analysis and Recommendations", lines=10)
+        ],
+        css=custom_css,
+        examples=load_examples_from_folder("./Examples"),
+        cache_examples=False,
+        article=model_explanation + pneumonia_info + disclaimer,
+        description=welcome_message,
+        title=logo_html,
+        theme=gr.themes.Soft()
+    )
+    
+    return iface
 
 # Launch the app
 if __name__ == "__main__":
-    iface.launch(debug=True)
+    demo = create_interface()
+    demo.launch(debug=True)