services/summarizer.py

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")