File size: 1,499 Bytes
54cd9a7 | 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 | import torch
from torchvision import models, transforms
from torchvision.io import read_image
from torch import nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
backbone_spatial = models.mobilenet_v3_large(weights='DEFAULT').features.to(device)
backbone_motion = models.mobilenet_v3_small(weights='DEFAULT').features.to(device)
backbone_spatial.eval()
backbone_motion.eval()
pool = nn.AdaptiveAvgPool2d(1)
norm = transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
resize = transforms.Resize((224, 224))
def extract_features(frame_paths):
imgs = []
for p in frame_paths:
img = read_image(p).float() / 255.0
imgs.append(norm(resize(img)))
imgs = torch.stack(imgs).to(device)
with torch.no_grad():
spatial = imgs[::3]
spatial_feat = pool(backbone_spatial(spatial)).flatten(1)
diffs = imgs[1:] - imgs[:-1]
diffs = torch.where(torch.abs(diffs) > 0.08, diffs, torch.zeros_like(diffs))
motion_feat = pool(backbone_motion(diffs)).flatten(1)
if spatial_feat.shape[0] < motion_feat.shape[0]:
pad = motion_feat.shape[0] - spatial_feat.shape[0]
spatial_feat = torch.cat([spatial_feat,
spatial_feat[-1:].repeat(pad, 1)])
else:
spatial_feat = spatial_feat[:motion_feat.shape[0]]
features = torch.cat([spatial_feat, motion_feat], dim=1)
return features.cpu()
|