Update app.py

app.py

@@ -13,13 +13,20 @@ target_labels = ['Pneumonia', 'Consolidation', 'Edema']
 target_idxs = [labels.index(lbl) for lbl in target_labels]
 
 def predict(image):
-image = image.convert("RGB").resize((224, 224))
+# Make sure image is RGB
+if image.mode != "RGB":
+image = image.convert("RGB")
+
+```
+# Process the image properly
 inputs = processor(images=image, return_tensors="pt")
+
 with torch.no_grad():
-logits = model(**inputs).logits
-probs = torch.sigmoid(logits).squeeze()
+    logits = model(**inputs).logits
+
+# Keep batch dimension for safety
+probs = torch.sigmoid(logits)[0]  # [batch, num_labels] -> [num_labels]
 
-```
 detected = []
 results = []
 for idx, lbl in zip(target_idxs, target_labels):