Update inference.py

inference.py

@@ -2,32 +2,34 @@ import torch
 from torchvision import transforms
 from PIL import Image
 import numpy as np
-from model import ConvLSTMModel, BetaVAE
+from model import BetaVAE, ConvLSTM  # your models
+import torch.nn.functional as F
 
 # ===============================================================
-# Config
+# Configuration
 # ===============================================================
 SEQUENCE_LENGTH = 10
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+
 # ===============================================================
-# Load Models
+# Model Loading
 # ===============================================================
 def load_convlstm(path="convlstm_model.pth"):
-    model = ConvLSTMModel()
+    model = ConvLSTM(input_channels=1, hidden_channels=[64, 64, 64], output_channels=1)
     checkpoint = torch.load(path, map_location=DEVICE)
     model.load_state_dict(checkpoint)
     model.eval().to(DEVICE)
-    print("✅ ConvLSTM model loaded.")
+    print("✅ ConvLSTM model loaded successfully.")
     return model
 
 
 def load_beta_vae(path="beta_vae_model.pth"):
-    model = BetaVAE()
+    model = BetaVAE(input_channels=1, latent_dim=64, beta=4.0)
     checkpoint = torch.load(path, map_location=DEVICE)
     model.load_state_dict(checkpoint)
     model.eval().to(DEVICE)
-    print("✅ β-VAE model loaded.")
+    print("✅ β-VAE model loaded successfully.")
     return model
 
 
@@ -35,7 +37,7 @@ def load_beta_vae(path="beta_vae_model.pth"):
 # Frame Pre/Post Processing
 # ===============================================================
 def preprocess_frame(frame: Image.Image):
-    """Convert PIL image → torch tensor (1,1,H,W) normalized to [0,1]."""
+    """Convert a PIL image to a normalized tensor (1, 1, 64, 64)."""
     transform = transforms.Compose([
         transforms.Grayscale(),
         transforms.Resize((64, 64)),
@@ -46,7 +48,7 @@ def preprocess_frame(frame: Image.Image):
 
 
 def postprocess_frame(tensor):
-    """Convert torch tensor (1,1,H,W) → PIL image."""
+    """Convert tensor (1, 1, H, W) → PIL image."""
     tensor = tensor.detach().cpu().clamp(0, 1)
     arr = tensor.squeeze().numpy() * 255
     arr = arr.astype(np.uint8)
@@ -54,33 +56,33 @@ def postprocess_frame(tensor):
 
 
 # ===============================================================
-# Inference Logic
+# Inference Helpers
 # ===============================================================
 @torch.no_grad()
-def predict_next_frame(model, sequence):
+def predict_next_frame(convlstm_model, sequence):
     """
+    Predict the next frame using the ConvLSTM model.
     Args:
-        model: ConvLSTMModel
+        convlstm_model: trained ConvLSTM
         sequence: tensor (1, T, 1, H, W)
     Returns:
-        PIL.Image
+        PIL.Image: predicted next frame
     """
-    model.eval()
     sequence = sequence.to(DEVICE)
-    next_frame = model(sequence)  # (1,1,H,W)
+    next_frame = convlstm_model(sequence)
     return postprocess_frame(next_frame)
 
 
 @torch.no_grad()
-def reconstruct_frame(model, frame):
+def reconstruct_frame(beta_vae_model, frame):
     """
+    Reconstruct a single frame using the β-VAE.
     Args:
-        model: BetaVAE
-        frame: tensor (1,1,H,W)
+        beta_vae_model: trained β-VAE
+        frame: tensor (1, 1, H, W)
     Returns:
-        PIL.Image
+        PIL.Image: reconstructed frame
     """
-    model.eval()
     frame = frame.to(DEVICE)
-    recon, mu, logvar = model(frame)
+    recon, mu, logvar = beta_vae_model(frame)
     return postprocess_frame(recon)