Update app.py

app.py

@@ -5,9 +5,7 @@ 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")
 
@@ -15,31 +13,24 @@ 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
+    return label_tuberculosis
 
 
 # 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"
+title = "Automated Classification of Tuberculosis using Machine Learning"
+description = "Upload your lungs Radiograph to find out if you are having Tuberculosis"
 
 css_code = ".gradio-container {background: url(https://actionmenshealth.com/wp-content/uploads/2018/05/heart.jpg); background-size: cover;}"