Views
No views yet
1import torch
2import torch.nn as nn
3
4
5class UCFModel(nn.Module):
6 def __init__(self, model_name="i3d_r50"):
7 super().__init__()
8 self.model_name = model_name
9
10 self.model = torch.hub.load("facebookresearch/pytorchvideo", model_name, pretrained=True)
11
12 in_features = self.model.blocks[-1].proj.in_features
13 self.model.blocks[-1].proj = nn.Linear(in_features, 2)
14
15 def forward(self, frames):
16 return self.model(frames)
171import torch
2from PIL import Image
3from huggingface_hub import hf_hub_download
4from torchvision import transforms
5
6
7inference_transform = transforms.Compose(
8 [
9 transforms.Resize((224, 224)),
10 transforms.ToTensor(),
11 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
12 ]
13)
14
15
16class UCFInferenceByFrames:
17 def __init__(self, repo_id):
18 self.repo_id = repo_id
19
20 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
21 self.model = self.load_model()
22
23 def load_model(self):
24 model_path = hf_hub_download(repo_id=self.repo_id, filename="ucf_model.pth")
25 state_dict = torch.load(model_path)
26
27 model = UCFModel().to(device=self.device)
28 model.load_state_dict(state_dict)
29 model.eval()
30
31 return model
32
33 def inference(self, frames):
34 video_tensor_list = []
35 for frame in frames:
36 frame_pil = Image.fromarray(frame)
37 frame_tensor = inference_transform(frame_pil)
38 video_tensor_list.append(frame_tensor)
39
40 video_tensor = torch.stack(video_tensor_list)
41 video_tensor = video_tensor.permute(1, 0, 2, 3).unsqueeze(0).float()
42
43 video_tensor = video_tensor.to(self.device)
44
45 with torch.no_grad():
46 output = self.model(video_tensor)
47
48 return output.argmax(1)1import cv2 as cv
2import numpy as np
3
4ucf = UCFInferenceByFrames("amjad-awad/ucf-i3d-model-by-3-block-lr-0.001")
5
6def inference(ucf_model, video_path, max_frames=16):
7 cap = cv.VideoCapture(video_path)
8
9 if not cap.isOpened():
10 print("No video")
11 return
12
13 frames = []
14
15 while True:
16 ret, frame = cap.read()
17
18 if not ret:
19 break
20
21 frames.append(frame)
22
23 length = len(frames)
24 indices = np.linspace(0, length - 1, max_frames, dtype=int)
25 frames = [frames[i] for i in indices]
26 predict = ucf_model.inference(frames)
27
28 return "Crime" if int(predict) == 1 else "No-Crime"
291predict = inference(ucf_model=ucf, video_path="YOUR_VIDEO_PATH.mp4")
2print(predict)