File size: 3,087 Bytes
3e586d9 d6f9e72 3e586d9 d6f9e72 3e586d9 312c423 3e586d9 95716e3 a55adf1 f9b41fc f335b44 f9b41fc ad69c6b fbb55a1 f9b41fc ad69c6b 312c423 fbb55a1 f9b41fc 312c423 fbb55a1 f9b41fc 0b02349 f9b41fc ad69c6b fbb55a1 f448117 f9b41fc f89e4ac f9b41fc 24a59af a55adf1 24a59af d57fde6 3e586d9 e829e37 d71bba9 a55adf1 6db0f0f a55adf1 22caeae b5e17d3 d71bba9 3e586d9 4da2443 d71bba9 24a59af 10a5c50 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import gradio as gr
from transformers import ViTForImageClassification, ViTFeatureExtractor
from PIL import Image
import torch
import numpy as np
# Load the pre-trained model and preprocessor (feature extractor)
model_name_pneumonia = "runaksh/chest_xray_pneumonia_detection"
model_name_tuberculosis = "runaksh/chest_xray_tuberculosis_detection"
model_pneumonia = ViTForImageClassification.from_pretrained(model_name_pneumonia)
model_tuberculosis = ViTForImageClassification.from_pretrained(model_name_tuberculosis)
feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
def classify_image(image):
# Convert the PIL Image to a format compatible with the feature extractor
image = np.array(image)
# Preprocess the image and prepare it for the model
inputs_pneumonia = feature_extractor(images=image, return_tensors="pt")
inputs_tuberculosis = feature_extractor(images=image, return_tensors="pt")
# Make prediction
with torch.no_grad():
outputs_pneumonia = model_pneumonia(**inputs_pneumonia)
logits_pneumonia = outputs_pneumonia.logits
outputs_tuberculosis = model_tuberculosis(**inputs_tuberculosis)
logits_tuberculosis = outputs_tuberculosis.logits
# Retrieve the highest probability class label index
predicted_class_idx_pneumonia = logits_pneumonia.argmax(-1).item()
predicted_class_idx_tuberculosis = logits_tuberculosis.argmax(-1).item()
# Define a manual mapping of label indices to human-readable labels
index_to_label_pneumonia = {0: "Pneumonia = NO",1: "Pneumonia = YES"}
index_to_label_tuberculosis = {0: "Tuberculosis = NO",1: "Tuberculosis = YES"}
# Convert the index to the model's class label
label_pneumonia = index_to_label_pneumonia.get(predicted_class_idx_pneumonia, "Unknown Label")
label_tuberculosis = index_to_label_tuberculosis.get(predicted_class_idx_tuberculosis, "Unknown Label")
label = label_pneumonia+".................."+label_tuberculosis
return label
# Create title, description and article strings
title = "Automated Classification of Pneumonia and Tuberculosis using Machine Learning"
description = "Upload your lungs Radiograph to find out if you are having Pneumonia or Tuberculosis"
css_code = ".gradio-container {background: url(https://www.bioworld.com/ext/resources/Stock-images/Therapeutic-topics/Respiratory/Respiratory-lungs-wireframe.png?1588285653); background-size: cover;}"
# Create Gradio interface
iface = gr.Interface(fn=classify_image,
inputs=gr.Image(), # Accepts image of any size
outputs=gr.Label(),
title=title,
description=description,
css=css_code
)
# Launch the app
iface.launch()
css_code = f"""
.gradio-container {{
background-image: url('{background_image_path}');
background-size: cover; /* Ensure image covers the container */
background-position: center; /* Center the image */
/* Add other styling options (e.g., padding, color) */
}}
"""
|