Update app.py

app.py

@@ -26,12 +26,12 @@ DEFAULT_PARAMS = {
         "test_seed": 42,   # must be non-negative
     },
     "image":{
-        "dataset_name": "QuotaClimat/frugalaichallenge-image-train",
+        "dataset_name": "pyronear/pyro-sdis",
         "test_size": 0.2,  # must be between 0 and 1
         "test_seed": 42,   # must be non-negative
     },
     "audio":{
-        "dataset_name": "QuotaClimat/frugalaichallenge-audio-train",
+        "dataset_name": "rfcx/frugalai",
         "test_size": 0.2,  # must be between 0 and 1
         "test_seed": 42,   # must be non-negative
     }
@@ -61,19 +61,33 @@ def evaluate_model(task: str, space_url: str):
         
         results = response.json()
         
-        # Check for required keys
-        required_keys = {
-            "username", "space_url", "submission_timestamp", "model_description", 
-            "accuracy", "energy_consumed_wh", "emissions_gco2eq", "emissions_data",
+        # Check for required keys based on task
+        base_required_keys = {
+            "username", "space_url", "submission_timestamp", "model_description",
+            "energy_consumed_wh", "emissions_gco2eq", "emissions_data",
             "api_route", "dataset_config"
         }
+
+        # Add task-specific accuracy keys
+        if task == "image":
+            accuracy_keys = {"classification_accuracy", "mean_iou"}
+        else:  # text and audio
+            accuracy_keys = {"accuracy"}
+        
+        required_keys = base_required_keys | accuracy_keys
         
         missing_keys = required_keys - set(results.keys())
         if missing_keys:
             return None, None, None, gr.Warning(f"API response missing required keys: {', '.join(missing_keys)}")
         
+        # Return appropriate accuracy metric based on task
+        if task == "image":
+            accuracy = results["classification_accuracy"]  # For display in UI
+        else:
+            accuracy = results["accuracy"]
+        
         return (
-            results["accuracy"],
+            accuracy,
             results["emissions_gco2eq"],
             results["energy_consumed_wh"],
             results