Update app.py

app.py

@@ -1,3 +1,48 @@
+Hugging Face's logo
+Hugging Face
+Search models, datasets, users...
+Models
+Datasets
+Spaces
+Posts
+Docs
+Solutions
+Pricing
+
+
+
+Spaces:
+
+runaksh
+/
+chest_xray_pneumonia_detection
+
+
+like
+0
+
+Logs
+App
+Files
+Community
+Settings
+chest_xray_pneumonia_detection
+/
+app.py
+
+runaksh's picture
+runaksh
+Update app.py
+3e586d9
+VERIFIED
+10 days ago
+raw
+history
+blame
+edit
+delete
+No virus
+1.58 kB
 import gradio as gr
 from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
@@ -5,75 +50,43 @@ import torch
 import numpy as np
 
 # Load the pre-trained model and preprocessor (feature extractor)
-model_name = "runaksh/chest_xray_pneumonia_detection"
-model = ViTForImageClassification.from_pretrained(model_name)
+model_name_pneumonia = "runaksh/chest_xray_pneumonia_detection"
+model_pneumonia = ViTForImageClassification.from_pretrained(model_name_pneumonia)
 feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
 
-def classify_image_pneumonia(image):
+def classify_image(image):
     # Convert the PIL Image to a format compatible with the feature extractor
-    image = np.array(image)
+    image_pneumonia = np.array(image)
     # Preprocess the image and prepare it for the model
-    inputs = feature_extractor(images=image, return_tensors="pt")
+    inputs_pneumonia = feature_extractor(images=image, return_tensors="pt")
     # Make prediction
     with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
+        outputs_pneumonia = model_pneumonia(**inputs_pneumonia)
+        logits_pneumonia = outputs.logits
     # Retrieve the highest probability class label index
-    predicted_class_idx = logits.argmax(-1).item()
+    predicted_class_idx_pneumonia = logits_pneumonia.argmax(-1).item()
     # Define a manual mapping of label indices to human-readable labels
-    index_to_label = {
+    index_to_label_pneumonia = {
         0: "NORMAL",
         1: "PNEUMONIA"
     }
 
     # Convert the index to the model's class label
-    label = index_to_label.get(predicted_class_idx, "Unknown Label")
+    label_pneumonia = index_to_label_pneumonia.get(predicted_class_idx_pneumonia, "Unknown Label")
 
-    return label
+    return label_pneumonia
 
-def classify_image_tuberculosis(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 = feature_extractor(images=image, return_tensors="pt")
-    # Make prediction
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-    # Retrieve the highest probability class label index
-    predicted_class_idx = logits.argmax(-1).item()
-    # Define a manual mapping of label indices to human-readable labels
-    index_to_label = {
-        0: "PNEUMONIA = NO",
-        1: "PNEUMONIA = YES"
-    }
+# Create title, description and article strings
+title = "Classification Demo"
+description = "XRay classification"
 
-    # Convert the index to the model's class label
-    label = index_to_label.get(predicted_class_idx, "Unknown Label")
+# Create Gradio interface
+iface = gr.Interface(fn=classify_image, 
+                     inputs=gr.Image(),  # Accepts image of any size
+                     outputs=gr.label_pneumonia(),
+                     title=title,
+                     description=description)
 
-    return label
+# Launch the app
+iface.launch()
 
-# Create Gradio interface
-def make_block(dem):
-    with dem:
-        gr.Markdown("Medical - Lungs Disease Prediction")
-        with gr.Tabs():
-            with gr.TabItem("Pneumonia Detection"):
-                with gr.Row():
-                    in_prompt_1 = gr.Image()
-                    out_response_1 = gr.Label()
-                b1 = gr.Button("Enter")
-
-            with gr.TabItem("Tuberculosis Detection"):
-                with gr.Row():
-                    in_prompt_2 = gr.Image()
-                    out_response_2 = gr.Label()
-                b2 = gr.Button("Enter")
-            b1.Block(classify_image_pneumonia, inputs=in_prompt_1, outputs=out_response_1)
-            b2.Block(classify_image_tuberculosis, inputs=in_prompt_2, outputs=out_response_2)
-           
-if __name__ == '__main__':
-
-    demo = gr.Blocks()
-    make_block(demo)
-    demo.launch()