Spaces:
Sleeping
Sleeping
| 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() | |