File size: 1,674 Bytes
d2542a3 6a6ee7b 7cd255e deda5ff d2542a3 6a6ee7b 0526375 d2542a3 1ff63b5 6a6ee7b 1ff63b5 d2542a3 7cd255e 1ff63b5 c5e3f79 7cd255e d2542a3 7cd255e deda5ff d2542a3 deda5ff d2542a3 deda5ff 7cd255e d2542a3 deda5ff d2542a3 deda5ff |
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 47 48 49 50 51 52 53 54 55 56 |
import cv2
import torch
from config import SCORE_THRESHOLD
from services.model_loader import load_model
import subprocess
import os
import numpy as np
from moviepy import ImageSequenceClip
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = load_model("Model/epoch-199.pkl")
model = model.to(device)
model = model.eval()
def get_scores(features):
# features.shape: (N, 1024)
# features.dtype: torch.float32
# features.device: cpu
with torch.no_grad():
print("Features before model inference:", features.shape)
scores, _ = model(features)
scores = scores.squeeze().cpu().numpy()
print("Features after model inference:", features.shape)
return scores
def get_selected_indices(scores, picks, threshold=SCORE_THRESHOLD):
selected = [picks[i] for i, score in enumerate(scores) if score >= threshold]
print("Threshold for selection:", threshold)
print("Scores:", len(scores))
print("Picks:", len(picks))
print("Selected:", len(selected), selected)
return selected
def save_summary_video(video_path, selected_indices, output_path, fps=15):
cap = cv2.VideoCapture(video_path)
selected = set(selected_indices)
frames = []
frame_id = 0
while True:
ret, frame = cap.read()
if not ret:
break
if frame_id in selected:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frames.append(frame)
frame_id += 1
cap.release()
if not frames:
print("No frames selected.")
return
clip = ImageSequenceClip(frames, fps=fps)
clip.write_videofile(output_path, codec="libx264")
|