|
import cv2 |
|
import torch |
|
from config import SCORE_THRESHOLD |
|
from services.model_loader import load_model |
|
|
|
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): |
|
|
|
with torch.no_grad(): |
|
scores, _ = model(features) |
|
return scores.squeeze().cpu().numpy() |
|
|
|
def get_selected_indices(scores, picks, threshold=SCORE_THRESHOLD): |
|
return [picks[i] for i, score in enumerate(scores) if score >= threshold] |
|
|
|
import subprocess |
|
import os |
|
|
|
def save_summary_video(video_path, selected_indices, output_path, fps=15): |
|
import cv2 |
|
|
|
cap = cv2.VideoCapture(video_path) |
|
selected = set(selected_indices) |
|
frame_id = 0 |
|
frames = {} |
|
|
|
while cap.isOpened(): |
|
ret, frame = cap.read() |
|
if not ret: |
|
break |
|
if frame_id in selected: |
|
frames[frame_id] = frame |
|
frame_id += 1 |
|
cap.release() |
|
|
|
if not frames: |
|
print("No frames selected.") |
|
return |
|
|
|
h, w, _ = list(frames.values())[0].shape |
|
|
|
|
|
raw_output_path = output_path.replace(".mp4", "_raw.mp4") |
|
writer = cv2.VideoWriter(raw_output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) |
|
for fid in sorted(frames.keys()): |
|
writer.write(frames[fid]) |
|
writer.release() |
|
|
|
|
|
try: |
|
subprocess.run([ |
|
"ffmpeg", |
|
"-y", |
|
"-i", raw_output_path, |
|
"-vcodec", "libx264", |
|
"-acodec", "aac", |
|
output_path |
|
], check=True) |
|
os.remove(raw_output_path) |
|
print(f"✅ FFmpeg re-encoded video saved to: {output_path}") |
|
except subprocess.CalledProcessError as e: |
|
print("❌ FFmpeg failed:", e) |
|
print("⚠️ Using raw video instead.") |
|
os.rename(raw_output_path, output_path) |
|
|