File size: 452 Bytes
c802d72
 
 
 
 
 
 
 
1022c4d
c802d72
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import torch
from assets.I3D_modelfiles.pytorch_i3d import InceptionI3d

class I3DModel:
    def __init__(self, weights_path, device):
        self.device = device

        self.model = InceptionI3d(400, in_channels=3)
        self.model.replace_logits(300)
        self.model.load_state_dict(torch.load(weights_path, map_location=device))
        self.model.to(device)
        self.model.eval()

    def __call__(self, x):
        return self.model(x)