Update app.py

app.py

@@ -16,24 +16,25 @@ def classify_image(image):
     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")
+    ####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
+        ####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()
+    ####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: "NORMAL",1: "PNEUMONIA"}
-    index_to_label_tuberculosis = {0: "NORMAL",1: "TUBERCULOSIS"}
+    ####index_to_label_tuberculosis = {0: "NORMAL",1: "TUBERCULOSIS"}
     # 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_tuberculosis = index_to_label_tuberculosis.get(predicted_class_idx_tuberculosis, "Unknown Label")
+    label = label_pneumonia
 
-    return label_pneumonia, label_tuberculosis
+    return label
 
 # Create title, description and article strings
 title = "Classification Demo"