Update app.py

app.py

@@ -54,24 +54,30 @@ def predict(image):
         # Make prediction
         output = model(image_tensor)
         prediction = output.argmax().item()
+        confidence = output.max().item()
         
+        # Print debugging information
+        print(f"Predicted: {prediction}, Confidence: {confidence}")
+        print(f"Model Output: {output}")
+
         # Class mapping
         class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
         
         # Return prediction result
-        return class_labels[prediction], output.max().item()
+        return class_labels[prediction], confidence
     except Exception as e:
         return {'error': str(e)}
 
+
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload an image"),
     outputs=[gr.Label(num_top_classes=3, label="Predicted Class"), gr.Number(label="Confidence")],
     examples=[
-        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer/default/train?p=2&row=201&image-viewer=8FBAF2C52C256A392660811C5659788734821C3A",
-        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer/default/train?p=2&image-viewer=AEF1AA2978EEB77362DA9CCC8792473666F7CDC6",
-        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer/default/train?image-viewer=C98E6CFAB26ECC3808C63185F6CCE90DE4E7C442"
+        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20290.jpg",
+        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20565.jpg",
+        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%208.jpg"
     ]
 )