Update app.py

app.py

@@ -1,29 +1,46 @@
-from transformers import ViTFeatureExtractor, ViTModel
+import gradio as gr
+from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
-import requests
+import torch
+import numpy as np
 
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
-loaded_model = ViTModel.from_pretrained("runaksh/chest_xray_pneumonia_detection")
-#inputs = feature_extractor(images=image, return_tensors="pt")
+# Load the pre-trained model and preprocessor (feature extractor)
+model_name = "runaksh/chest_xray_pneumonia_detection"
+model = ViTForImageClassification.from_pretrained(model_name)
+feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
 
-def predict(img):
-    #inputs = feature_extractor(images=image, return_tensors="pt")
-    pipe = pipeline('image-classification', model=model_name, device=0)
-    pred = pipe(image)
-    return pred
+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 = 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: "NORMAL",
+        1: "PNEUMONIA"
+    }
+
+    # Convert the index to the model's class label
+    label = index_to_label.get(predicted_class_idx, "Unknown Label")
+
+    return label
 
 # Create title, description and article strings
 title = "Classification Demo"
 description = "XRay classification"
 
-# Create the Gradio demo
-demo = gr.Interface(fn=predict, # mapping function from input to output
-                    inputs=gr.Image(type='filepath'), # what are the inputs?
-                    outputs=[gr.Label(label="Predictions"), # what are the outputs?
-                             gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
-                    examples=example_list,
-                    title=title,
-                    description=description,)
+# Create Gradio interface
+iface = gr.Interface(fn=classify_image, 
+                     inputs=gr.Image(),  # Accepts image of any size
+                     outputs=gr.Label(),
+                     title=title,
+                     description=description)
 
-# Launch the demo!
-demo.launch()
+# Launch the app
+iface.launch()