PneumonoBot

PneumonoBot - Pneumonia Classification and Chatbot

PneumonoBot is an advanced system that leverages a fine-tuned Vision Transformer (ViT) model to accurately classify chest X-ray images as either 'Normal' or 'Pneumonia'. In addition, it features an intelligent chatbot that can answer various questions related to pneumonia.

GitHub Repository

You can explore the complete details of PneumonoBot, including the model architecture, training procedures, and chatbot implementation, in the official GitHub repository: PneumonoBot Repository.

🚀 Save the Model and Feature Extractor Locally

To ensure fast and efficient processing, you can save the model and its feature extractor directly to your local machine. This allows for easy offline usage and quicker model loading.

from transformers import ViTForImageClassification, ViTFeatureExtractor

# Define the model path from Hugging Face
model_name = "M-Yaqoob/PneumonoBot"

# Load the pre-trained model and feature extractor
model = ViTForImageClassification.from_pretrained(model_name)
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)

# Specify the directory to save the model and feature extractor
save_directory = "./vit_classification_pneumonobot"

# Save both locally for future use
model.save_pretrained(save_directory)
feature_extractor.save_pretrained(save_directory)

print(f"Model and feature extractor saved to {save_directory}!")

🧠 Load the Model and Feature Extractor from Local Storage

Need to run predictions? Easily load your locally saved model and feature extractor with the following code:

from transformers import ViTForImageClassification, ViTFeatureExtractor
import torch
from PIL import Image
import torchvision.transforms as transforms

# Set the path where the model is stored
save_directory = "./vit_classification_pneumonobot"

# Load the model and feature extractor from the local directory
model = ViTForImageClassification.from_pretrained(save_directory)
feature_extractor = ViTFeatureExtractor.from_pretrained(save_directory)

print("Model and feature extractor successfully loaded!")

🖼️ Provide an Image for Prediction

Now, let’s put the model to work by classifying an X-ray image as either 'Normal' or 'Pneumonia'. The steps below guide you through image loading, preprocessing, and prediction.

# Define the image path
# Example for Pneumonia image
# image_path = "./Images/PNEUMONIA/person66_virus_125.jpeg" 

# Example for Normal image
image_path = "./Images/NORMAL/IM-0069-0001.jpeg"

# Open and preprocess the image
image = Image.open(image_path).convert("RGB")
inputs = feature_extractor(images=image)

# Prepare inputs for the model
inputs = {key: torch.tensor(value) for key, value in inputs.items()}

# Manually define the label mapping
labels = {0: 'Normal', 1: 'Pneumonia'}

# Perform inference with no gradient calculation
with torch.no_grad():
    outputs = model(**inputs)

# Get the index of the predicted class and map it to the label
predicted_class_idx = outputs.logits.argmax(-1).item()
predicted_label = labels[predicted_class_idx]

print(f"🔍 Predicted label: {predicted_label}")

🎯 Result:

With this workflow, you can confidently predict whether an X-ray indicates a normal lung condition or pneumonia, powered by state-of-the-art Vision Transformer (ViT) models.

Model Accuracy

By fine-tuning the Vision Transformer (ViT) model, we achieved an impressive 91% accuracy on the test dataset.