File size: 2,211 Bytes
a14c800 | 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 | """Predict next frame using improved U-Net v2 (fp16 weights)."""
import json
import sys
from typing import Any
import numpy as np
import torch
sys.path.insert(0, '/home/coder/code')
from unet_v3 import UNetV3
def load_model(model_dir: str) -> Any:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
with open(f"{model_dir}/config.json", 'r') as f:
config = json.load(f)
model = UNetV3(
in_channels=config['in_channels'],
out_channels=config['out_channels'],
enc_channels=config['enc_channels'],
dec_channels=config['dec_channels'],
use_tanh=config.get('use_tanh', False),
)
state_fp16 = torch.load(f"{model_dir}/model.pt", map_location='cpu', weights_only=True)
state_fp32 = {k: v.float() for k, v in state_fp16.items()}
model.load_state_dict(state_fp32)
model.to(device)
model.eval()
return {'model': model, 'device': device, 'n_context': config['n_context']}
def predict_next_frame(model_dict: Any, context_frames: np.ndarray) -> np.ndarray:
model = model_dict['model']
device = model_dict['device']
n_ctx = model_dict['n_context']
# Use last n_ctx frames, pad if needed
if len(context_frames) >= n_ctx:
frames = context_frames[-n_ctx:]
else:
pad_count = n_ctx - len(context_frames)
padding = np.stack([context_frames[0]] * pad_count, axis=0)
frames = np.concatenate([padding, context_frames], axis=0)
# Normalize to [0, 1], convert to CHW
frames_norm = frames.astype(np.float32) / 255.0
last_frame = frames_norm[-1]
# Stack context: (n_ctx, 3, H, W) -> (n_ctx*3, H, W)
stacked = np.transpose(frames_norm, (0, 3, 1, 2)).reshape(-1, 64, 64)
last_chw = np.transpose(last_frame, (2, 0, 1))
with torch.no_grad():
inp = torch.from_numpy(stacked).unsqueeze(0).to(device)
last_t = torch.from_numpy(last_chw).unsqueeze(0).to(device)
delta = model(inp)
pred = torch.clamp(last_t + delta, 0, 1)
pred = pred.squeeze(0).cpu().numpy()
# Convert to uint8 HWC
pred = np.clip(pred, 0, 1)
pred = (pred * 255.0).astype(np.uint8)
return np.transpose(pred, (1, 2, 0))
|