Spaces:
Sleeping
Sleeping
File size: 3,598 Bytes
7ad2c42 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 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 |