Spaces:
Sleeping
Sleeping
| import os | |
| import gradio as gr | |
| from PIL import Image | |
| import torch | |
| from transformers import ViTForImageClassification, ViTImageProcessor | |
| from datasets import load_dataset, DownloadConfig | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import cv2 | |
| import requests | |
| # Mistral AI API configuration | |
| MISTRAL_API_KEY = "eoiBrPQzLjwNgOFgD7I4A4XF3TJOgBet" | |
| MISTRAL_API_URL = "https://api.mistral.ai/v1/chat/completions" | |
| def get_mistral_completion(prompt): | |
| headers = { | |
| "Content-Type": "application/json", | |
| "Authorization": f"Bearer {MISTRAL_API_KEY}" | |
| } | |
| data = { | |
| "model": "mistral-medium", | |
| "messages": [{"role": "user", "content": prompt}], | |
| "temperature": 0.7 | |
| } | |
| response = requests.post(MISTRAL_API_URL, headers=headers, json=data) | |
| if response.status_code == 200: | |
| return response.json()['choices'][0]['message']['content'] | |
| else: | |
| return "Error getting AI explanation. Please try again." | |
| # Model and processor configuration | |
| model_name_or_path = "google/vit-base-patch16-224-in21k" | |
| processor = ViTImageProcessor.from_pretrained(model_name_or_path) | |
| # 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) | |
| class_names = train_dataset.features["label"].names | |
| # Load ViT model | |
| model = ViTForImageClassification.from_pretrained( | |
| "./models", | |
| num_labels=len(class_names), | |
| id2label={str(i): label for i, label in enumerate(class_names)}, | |
| label2id={label: i for i, label in enumerate(class_names)}, | |
| ) | |
| model.eval() | |
| 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, please remember that NORMAL ['normal'] means this user does not have any Pneumonia. | |
| 2. Suggest when they should still consider consulting a doctor even though this user does not have any penumonia as per the test result. | |
| 3. List key symptoms that would warrant medical attention. Always identify yourself as PneumoInsight Bot. | |
| 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 . Always identify yourself as PneumoInsight Bot. | |
| 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.""" | |
| return get_mistral_completion(prompt) | |
| 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 | |
| ) | |
| return {"probabilities": result, "heatmap": heatmap_img} | |
| 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 | |
| def load_examples_from_folder(folder_path): | |
| examples = [] | |
| if os.path.exists(folder_path): | |
| for file in os.listdir(folder_path): | |
| if file.endswith((".png", ".jpg", ".jpeg")): | |
| examples.append(os.path.join(folder_path, file)) | |
| return examples | |
| 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__": | |
| demo = create_interface() | |
| demo.launch(debug=True) |