Upload: load.py

single/load.py

@@ -14,8 +14,8 @@ class EnsembleModel(torch.nn.Module):
         super(EnsembleModel, self).__init__()
         self.models = torch.nn.ModuleList(models)
         self.mode = mode
-        if mode not in ["min", "mean", "max", "none"]:
-            raise ValueError("Mode must be 'none', 'min', 'mean', or 'max'.")
+        if mode not in ["min", "mean", "median", "max", "none"]:
+            raise ValueError("Mode must be 'none', 'min', 'mean', 'median', or 'max'.")
 
     def forward(self, x) -> Tuple[torch.Tensor, torch.Tensor]:
         """
@@ -38,18 +38,20 @@ class EnsembleModel(torch.nn.Module):
         stacked_outputs = torch.stack(outputs, dim=1)  # (B, N, 1, H, W)
         stacked_outputs = stacked_outputs.squeeze(2)   # (B, N, H, W)
         
-        # Calculate mean (probabilities)
+        # Calculate aggregated probabilities
         if self.mode == "max":
             output_probs = torch.max(stacked_outputs, dim=1, keepdim=True)[0]
         elif self.mode == "mean":
             output_probs = torch.mean(stacked_outputs, dim=1, keepdim=True)
+        elif self.mode == "median":
+            output_probs = torch.median(stacked_outputs, dim=1, keepdim=True)[0]
         elif self.mode == "min":
             output_probs = torch.min(stacked_outputs, dim=1, keepdim=True)[0]
         elif self.mode == "none":
             # Return all predictions without aggregation
             return stacked_outputs, None
         else:
-            raise ValueError("Mode must be 'min', 'mean', or 'max'.")
+            raise ValueError("Mode must be 'min', 'mean', 'median', or 'max'.")
         
         # Calculate uncertainty (normalized standard deviation)
         N = len(outputs)
@@ -69,16 +71,19 @@ class EnsembleModel(torch.nn.Module):
             uncertainty = torch.zeros_like(output_probs)
         
         return output_probs, uncertainty  # Both (B, 1, H, W)
-
+    
 def compiled_model(
     path: pathlib.Path, 
     stac_item: pystac.Item,
-    mode: Literal["min", "mean", "max"] = "max",
+    mode: Literal["min", "mean", "median", "max"] = "max",
     *args, **kwargs
 ) -> EnsembleModel:
     """
     Loads the ensemble dynamically using the 'stac_item' as the source of truth.
     
+    Args:
+        mode: Aggregation mode - 'min', 'mean', 'median', or 'max' (default: 'max')
+    
     Returns an EnsembleModel that outputs (probabilities, uncertainty) tuple.
     """
     model_paths = []