Create amlm.py

core/amlm.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import os
+
+# ==========================================
+# 1. MODEL CONFIGURATION (V7 Standard)
+# ==========================================
+# Fixed latent dimension matching your trained "amlm_v5.pth"
+LATENT_DIM = 64
+
+class AMLM_3D(nn.Module):
+    def __init__(self):
+        super(AMLM_3D, self).__init__()
+        
+        # --- ENCODER ---
+        # Input: (Batch, 6, 64, 9) 
+        # 6 Channels = 3 Value Channels (Data) + 3 Visibility Channels (Mask)
+        self.encoder = nn.Sequential(
+            # Layer 1: 64x9 -> 32x9
+            nn.Conv2d(6, 32, kernel_size=3, padding=1), 
+            nn.ReLU(), 
+            nn.MaxPool2d((2, 1)), 
+            
+            # Layer 2: 32x9 -> 16x9
+            nn.Conv2d(32, 64, kernel_size=3, padding=1), 
+            nn.ReLU(), 
+            nn.MaxPool2d((2, 1)), 
+            
+            # Flatten: 64 channels * 16 height * 9 width = 9216
+            nn.Flatten() 
+        )
+        
+        # --- BOTTLENECK ---
+        # Compresses the 9216 features into the 64-dim "Clinical Syndrome" vector
+        self.bottleneck = nn.Sequential(
+            nn.Linear(9216, 256), 
+            nn.ReLU(), 
+            nn.Linear(256, LATENT_DIM) 
+        )
+        
+        # --- DECODER INPUT ---
+        # Expands 64-dim syndrome back to spatial features
+        self.decoder_input = nn.Sequential(
+            nn.Linear(LATENT_DIM, 256), 
+            nn.ReLU(), 
+            nn.Linear(256, 9216), 
+            nn.ReLU()
+        )
+        
+        # --- DECODER ---
+        # Reconstructs the full probability map
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(64, 32, kernel_size=3, padding=1), 
+            nn.Upsample(scale_factor=(2, 1)), 
+            nn.ReLU(),
+            
+            nn.ConvTranspose2d(32, 16, kernel_size=3, padding=1), 
+            nn.Upsample(scale_factor=(2, 1)), 
+            nn.ReLU(),
+            
+            # Output: 3 Channels (Initial, Medial, Final)
+            nn.Conv2d(16, 3, kernel_size=3, padding=1), 
+            nn.Sigmoid() # Squish output to 0.0 - 1.0 probability
+        )
+
+    def forward(self, x):
+        """
+        Forward pass for inference.
+        x shape must be: (Batch_Size, 6, 64, 9)
+        """
+        features = self.encoder(x)
+        latent = self.bottleneck(features) 
+        
+        x_recon = self.decoder_input(latent)
+        # Reshape back to spatial dimensions (Batch, Channels, Height, Width)
+        x_recon = x_recon.view(-1, 64, 16, 9) 
+        
+        # Output map (Batch, 3, 64, 9)
+        output = self.decoder(x_recon)
+        
+        return output, latent
+
+# ==========================================
+# 2. LOADER UTILITY
+# ==========================================
+def load_amlm_model(model_path, device):
+    """
+    Initializes the model and loads weights safely.
+    """
+    print(f"--- Loading AMLM V7 from {model_path} ---")
+    model = AMLM_3D().to(device)
+    
+    if os.path.exists(model_path):
+        try:
+            state_dict = torch.load(model_path, map_location=device)
+            model.load_state_dict(state_dict)
+            model.eval() # Set to inference mode (freezes Dropout/BatchNorm)
+            print("✅ AMLM Weights Loaded Successfully")
+        except Exception as e:
+            print(f"❌ CRITICAL: Failed to load AMLM weights: {e}")
+            # We assume the app should crash or handle this if weights are bad
+            raise e
+    else:
+        print(f"⚠️ WARNING: Model file not found at {model_path}. Running with random weights (Debugging Mode).")
+    
+    return model