Spaces:
Sleeping
Sleeping
| 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 |