Views
No views yet
1import torch
2import torch.nn as nn
3import torchvision.models.video as video_models
4import cv2
5import numpy as np
6
7# Define model architecture
8class VideoClassifier(nn.Module):
9 def __init__(self, num_classes=6):
10 super(VideoClassifier, self).__init__()
11 self.backbone = video_models.r3d_18(pretrained=False)
12 in_features = self.backbone.fc.in_features
13 self.backbone.fc = nn.Linear(in_features, num_classes)
14
15 def forward(self, x):
16 return self.backbone(x)
17
18# Load model
19device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
20model = VideoClassifier(num_classes=6)
21model.load_state_dict(torch.load('best_model.pth', map_location=device))
22model.to(device)
23model.eval()
24
25# Preprocess video
26def preprocess_video(video_path, num_frames=90, target_size=(224, 224)):
27 cap = cv2.VideoCapture(video_path)
28 frames = []
29
30 while True:
31 ret, frame = cap.read()
32 if not ret:
33 break
34 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
35 frame = cv2.resize(frame, target_size)
36 frame = frame.astype(np.float32) / 255.0
37 frames.append(frame)
38
39 cap.release()
40
41 # Sample to 90 frames
42 if len(frames) < num_frames:
43 repeat_factor = int(np.ceil(num_frames / len(frames)))
44 frames = (frames * repeat_factor)[:num_frames]
45 elif len(frames) > num_frames:
46 indices = np.linspace(0, len(frames) - 1, num_frames).astype(int)
47 frames = [frames[i] for i in indices]
48
49 # Convert to tensor (C, T, H, W)
50 frames = torch.FloatTensor(np.array(frames)).permute(3, 0, 1, 2)
51
52 # Apply normalization
53 mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1, 1)
54 std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1, 1)
55 frames = (frames - mean) / std
56
57 # Per-video standardization
58 video_mean = frames.mean()
59 video_std = frames.std()
60 if video_std > 0:
61 frames = (frames - video_mean) / video_std
62
63 return frames.unsqueeze(0) # Add batch dimension
64
65# Inference
66video_tensor = preprocess_video("path/to/video.mp4")
67video_tensor = video_tensor.to(device)
68
69with torch.no_grad():
70 outputs = model(video_tensor)
71 probabilities = torch.softmax(outputs, dim=1)
72 predicted_class = torch.argmax(probabilities, dim=1).item()
73
74class_names = ["0%", "60%", "70%", "80%", "85%", "90%"]
75print(f"Predicted normality rate: {class_names[predicted_class]}")
76print(f"Confidence: {probabilities[0][predicted_class].item():.4f}")@misc{sperm-normality-classifier,
author = {Raid Athmane Benlala},
title = {Sperm Normality Rate Classifier},
year = {2025},
publisher = {Hugging Face},
howpublished = {\url{https://huggingface.co/raidAthmaneBenlala/normality-rate-classifier}}
}