Update app.py

app.py

@@ -46,7 +46,13 @@ model.to(device)
 def sigmoid(_outputs):
     return 1.0 / (1.0 + np.exp(-_outputs))
 
-def predict(image, threshold=0.5):
+def download_thresholds(repo_id, filename):
+    threshold_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    with open(threshold_path, 'r') as threshold_file:
+        thresholds = json.load(threshold_file)
+    return thresholds
+
+def predict(image, slider_threshold=0.5, fixed_thresholds=None):
     # Preprocess the image
     processor = AutoImageProcessor.from_pretrained(checkpoint_name)
     inputs = processor(images=image, return_tensors="pt").to(device)
@@ -57,13 +63,26 @@ def predict(image, threshold=0.5):
     logits = model_outputs.logits[0]
     probabilities = torch.sigmoid(logits).cpu().numpy()  # Convert to probabilities
 
-    # Create a dictionary of label scores
-    results = {id2label[str(i)]: float(prob) for i, prob in enumerate(probabilities)}
+    # Create a dictionary of label scores based on the slider threshold
+    slider_results = {id2label[str(i)]: float(prob) for i, prob in enumerate(probabilities) if prob > slider_threshold}
+    
+    # If fixed thresholds are provided, create a dictionary of label scores based on the fixed thresholds
+    fixed_threshold_results = None
+    if fixed_thresholds is not None:
+        fixed_threshold_results = {id2label[str(i)]: float(prob) for i, prob in enumerate(probabilities) if prob > fixed_thresholds[id2label[str(i)]]}
+
+    return slider_results, fixed_threshold_results
 
-    # Filter out predictions below a certain threshold (e.g., 0.5)
-    filtered_results = {label: prob for label, prob in results.items() if prob > threshold}
+def predict_wrapper(image, slider_threshold):
+    # Download thresholds from the model repository
+    thresholds = download_thresholds(checkpoint_name, "threshold.json")
+    
+    # Get predictions from the predict function using both the slider and fixed thresholds
+    slider_results, fixed_threshold_results = predict(image, slider_threshold, thresholds)
+    
+    # Return both sets of predictions for Gradio outputs
+    return slider_results, fixed_threshold_results
 
-    return filtered_results 
 
 # Define style
 title = "Victor - DinoVd'eau image classification"
@@ -71,11 +90,14 @@ model_link = "https://huggingface.co/" + checkpoint_name
 description = f"This application showcases the capability of artificial intelligence-based systems to identify objects within underwater images. To utilize it, you can either upload your own image or select one of the provided examples for analysis.\nFor predictions, we use this [open-source model]({model_link})"
 
 iface = gr.Interface(
-    fn=predict,
+    fn=predict_wrapper,
     inputs=[gr.components.Image(type="pil"), gr.components.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Threshold")],
-    outputs=gr.components.Label(),
+    outputs=[
+        gr.components.Label(label="Slider Threshold Predictions"),
+        gr.components.Label(label="Fixed Thresholds Predictions")
+    ],
     title=title, 
-    description = description,
+    description=description,
     examples=[["session_GOPR0106.JPG"], 
               ["session_2021_08_30_Mayotte_10_image_00066.jpg"], 
               ["session_2018_11_17_kite_Le_Morne_Manawa_G0065777.JPG"],