Update app.py

app.py

@@ -9,7 +9,29 @@ model_name = "runaksh/chest_xray_pneumonia_detection"
 model = ViTForImageClassification.from_pretrained(model_name)
 feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
 
-def classify_image(image):
+def classify_image_pneumonia(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
+
+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
@@ -47,8 +69,8 @@ def make_block(dem):
                     in_prompt_2 = gr.Image()
                     out_response_2 = gr.Label()
                 b2 = gr.Button("Enter")
-            b1.click(predict_sentiment, inputs=in_prompt_1, outputs=out_response_1)
-            b2.click(predict, inputs=in_prompt_2, outputs=out_response_2)
+            b1.click(classify_image_pneumonia, inputs=in_prompt_1, outputs=out_response_1)
+            b2.click(classify_image_tuberculosis, inputs=in_prompt_2, outputs=out_response_2)
            
 if __name__ == '__main__':