File size: 2,383 Bytes
5a6434e 7cf6bed 5a6434e 87bfad6 7cf6bed 87bfad6 5a6434e f5366ce 7cf6bed 87bfad6 5a6434e 87bfad6 7cf6bed 87bfad6 2fd1c51 87bfad6 5a6434e 87bfad6 2fd1c51 339daa1 87bfad6 5a6434e d4efc46 87bfad6 07239aa 5a6434e 87bfad6 f2eecc0 2fd1c51 87bfad6 f2eecc0 87bfad6 f2eecc0 87bfad6 f2eecc0 87bfad6 5a6434e 87bfad6 5a6434e 2e7cf8e 87bfad6 2e7cf8e 87bfad6 | 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 | """Prediction interface for Multi-Scale Flow-Warp-Mask U-Net v10 with TTA."""
import sys
import os
import numpy as np
import torch
sys.path.insert(0, "/home/coder/code")
from multiscale_flow_model import MultiScaleFlowUNet
from flownet_model import differentiable_warp
CONTEXT_LEN = 4
CHANNELS = [56, 112, 224]
def load_model(model_dir: str):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MultiScaleFlowUNet(in_channels=12, channels=CHANNELS)
model_path = os.path.join(model_dir, "model.pt")
state_dict = torch.load(model_path, map_location=device, weights_only=True)
state_dict = {k: v.float() for k, v in state_dict.items()}
model.load_state_dict(state_dict)
model.to(device)
model.eval()
return {"model": model, "device": device}
def _prepare_input(context_frames):
if len(context_frames) >= CONTEXT_LEN:
frames = context_frames[-CONTEXT_LEN:]
else:
pad_count = CONTEXT_LEN - len(context_frames)
padding = np.stack([context_frames[0]] * pad_count, axis=0)
frames = np.concatenate([padding, context_frames], axis=0)
frames_t = torch.from_numpy(frames.astype(np.float32) / 255.0)
frames_t = frames_t.permute(0, 3, 1, 2)
return frames_t
def _run_model(model, frames_t, device):
last_frame = frames_t[-1].unsqueeze(0)
inp = frames_t.reshape(1, -1, 64, 64)
inp = inp.to(device)
last_frame = last_frame.to(device)
flows, mask, gen_frame = model(inp)
# Use finest flow (last element)
flow = flows[-1]
warped = differentiable_warp(last_frame, flow)
pred = mask * warped + (1 - mask) * gen_frame
pred = torch.clamp(pred, 0, 1)
return pred
def predict_next_frame(model_dict, context_frames: np.ndarray) -> np.ndarray:
model = model_dict["model"]
device = model_dict["device"]
frames_t = _prepare_input(context_frames)
with torch.no_grad():
# Original prediction
pred1 = _run_model(model, frames_t, device)
# TTA: horizontally flipped prediction
frames_flipped = frames_t.flip(-1)
pred2_flipped = _run_model(model, frames_flipped, device)
pred2 = pred2_flipped.flip(-1)
# Average
pred = (pred1 + pred2) / 2.0
pred = pred[0].cpu().permute(1, 2, 0).numpy()
pred = (pred * 255).clip(0, 255).astype(np.uint8)
return pred
|