Spaces:
Sleeping
Sleeping
File size: 1,623 Bytes
98a77e0 |
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 |
from numpy.lib.npyio import load
from torch._C import device
import sys
sys.path.append('/scratch/shared/beegfs/szwu/projects/video3d/RAFT')
from core.raft import RAFT
from .utils import InputPadder
import torch
class AttrDict(dict):
def __init__(self, *args, **kwargs):
super(AttrDict, self).__init__(*args, **kwargs)
self.__dict__ = self
class FlowModel():
def __init__(self, model, device):
args = AttrDict({'model': model, 'small': False, 'mixed_precision': False, 'alternate_corr': False})
self.model = self.load_model(args, device)
self.device = device
@staticmethod
def load_model(args, device):
model = torch.nn.DataParallel(RAFT(args))
model.load_state_dict(torch.load(args.model))
model = model.module
model.to(device)
model.eval()
return model
def preprocess_image(self, image):
# image = image[:, :, ::-1].copy()
image = torch.from_numpy(image).permute(2, 0, 1).float()
image = image.to(self.device)
image = image[None]
# size = [540, 960]
# image = torch.nn.functional.interpolate(image, size=size, mode='bilinear', align_corners=False)
padder = InputPadder(image.shape)
return padder.pad(image)[0], padder
def compute_flow(self, frame, next_frame, iters=20):
frame, padder = self.preprocess_image(frame)
next_frame, padder = self.preprocess_image(next_frame)
_, flow = self.model(frame, next_frame, iters=iters, test_mode=True)
return padder.unpad(flow)[0].permute(1, 2, 0).cpu().numpy()
|