Upload HMS_EXP_4_MODEL.py with huggingface_hub

HMS_EXP_4_MODEL.py

@@ -0,0 +1,56 @@
+class CustomModel(nn.Module):
+    def __init__(self, config, num_classes: int = 6, pretrained: bool = True):
+        super(CustomModel, self).__init__()
+        self.USE_KAGGLE_SPECTROGRAMS = True
+        self.USE_EEG_SPECTROGRAMS    = True
+        self.model                   = timm.create_model(
+                                                            config.MODEL_NAME,
+                                                            pretrained=pretrained,
+                                                        )
+        if config.FREEZE:
+            for i,(name, param) in enumerate(list(self.model.named_parameters())[0:config.NUM_FROZEN_LAYERS]):
+                param.requires_grad = False
+
+        self.features      = nn.Sequential(*list(self.model.children())[:-2])
+        self.custom_layers = nn.Sequential(
+                                                nn.AdaptiveAvgPool2d(1),
+                                                nn.Flatten(),
+                                                nn.Linear(self.model.num_features, num_classes)
+                                            )
+
+    def __reshape_input(self, x):
+        """
+        Reshapes input (128, 256, 8) -> (786, 786, 3) monotone image.
+        """ 
+        # === Get spectograms ===
+        spectograms = [x[:, :, :, i:i+1] for i in range(4)]
+        spectograms = torch.cat(spectograms, dim=1)
+
+        # === Get EEG spectograms ===
+        eegs = [x[:, :, :, i:i+1] for i in range(4,8)]
+        eegs = torch.cat(eegs, dim=1)
+
+        # === Reshape (786, 786, 3) ===
+        if self.USE_KAGGLE_SPECTROGRAMS & self.USE_EEG_SPECTROGRAMS:
+            # Concatenate spectograms and eegs along the channels (dim=2)
+            x = torch.cat([spectograms, eegs], dim=2)
+        elif self.USE_EEG_SPECTROGRAMS:
+            x = eegs
+        else:
+            x = spectograms
+
+        # Replicate the single-channel data to create a monotone image
+        x = torch.cat([x, x, x], dim=3)
+
+        # Permute dimensions to match the desired shape (batch_size, channels, height, width)
+        x = x.permute(0, 3, 1, 2)
+
+        return x
+
+
+    
+    def forward(self, x):
+        x = self.__reshape_input(x)
+        x = self.features(x)
+        x = self.custom_layers(x)
+        return x