Update ensemble/load.py

ensemble/load.py

@@ -17,14 +17,15 @@ def load_model_module(model_path: pathlib.Path):
     return model
 
 class EnsembleModel(torch.nn.Module):
-    def __init__(self, model1, model2, model3, model4, model5, mode="max"):
+    def __init__(self, model1, model2, model3, model4, model5, model6, mode="max"):
         super(EnsembleModel, self).__init__()
         self.model1 = model1
         self.model2 = model2
         self.model3 = model3
         self.model4 = model4
         self.model5 = model5
-        self.models = [model1, model2, model3, model4, model5]
+        self.model6 = model6
+        self.models = [model1, model2, model3, model4, model5, model6]
         self.mode = mode
         if mode not in ["min", "mean", "max"]:
             raise ValueError("Mode must be 'min', 'mean', or 'max'.")
@@ -42,6 +43,8 @@ class EnsembleModel(torch.nn.Module):
             output_probs = torch.mean(torch.cat(outputs, dim=1), dim=1)[0].squeeze()
         elif self.mode == "min":
             output_probs = torch.min(torch.cat(outputs, dim=1), dim=1)[0].squeeze()
+        elif self.mode == "none":
+            return torch.cat(outputs, dim=1)
         else:
             raise ValueError("Mode must be 'min', 'mean', or 'max'.")
 
@@ -68,6 +71,7 @@ def compiled_model(path, device: str = "cpu", mode: Literal["min", "mean", "max"
     model3_f = path / "1dpwunetpp.safetensor"
     model4_f = path / "unet.safetensor"
     model5_f = path / "unetpp.safetensor"
+    model6_f = path / "c2r1km.safetensor"
 
     # Load model parameters
     model1_weights = safetensors.torch.load_file(model1_f)
@@ -75,6 +79,7 @@ def compiled_model(path, device: str = "cpu", mode: Literal["min", "mean", "max"
     model3_weights = safetensors.torch.load_file(model3_f)
     model4_weights = safetensors.torch.load_file(model4_f)
     model5_weights = safetensors.torch.load_file(model5_f)
+    model6_weights = safetensors.torch.load_file(model6_f)
 
     # Load all models
 
@@ -128,12 +133,29 @@ def compiled_model(path, device: str = "cpu", mode: Literal["min", "mean", "max"
     for param in model5.parameters():
         param.requires_grad = False
     model5 = model5.eval()
+
+    # Model 6 (C2R1KM)
+    model6 = load_model_module(path / "c2r1km.py").CloudMaskOne(
+        hidden_layer_sizes=(21, 20),
+        activation='relu',
+        last_activation='sigmoid',
+        dropout_rate=0.1,
+        input_dim=40,
+        batch_norm=False
+    )
+    model6.load_state_dict(model6_weights)
+    model6 = model6.to(device)
+    for param in model6.parameters():
+        param.requires_grad = False
+    model6 = model6.eval()
     
     # Create ensemble model
-    cloud_model = EnsembleModel(model1, model2, model3, model4, model5, mode=mode)
+    cloud_model = EnsembleModel(model1, model2, model3, model4, model5, model6, mode=mode)
 
     return cloud_model
 
+
+
 def display_results(path: pathlib.Path, device: str = "cpu", mode: Literal["min", "mean", "max"] ="max", *args, **kwargs):
     # Load model
     model = compiled_model(path, device, mode=mode)