src/utils.py

import os
import cv2
import time
import shutil
import logging
import uuid
import torch
import numpy as np
import atexit
from concurrent.futures import ThreadPoolExecutor
from typing import Union

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

thread_pool_executor = ThreadPoolExecutor(max_workers=2)

def delete_later(path: Union[str, os.PathLike], delay: int = 600):
    def _delete():
        try:
            if os.path.isfile(path):
                os.remove(path)
            elif os.path.isdir(path):
                shutil.rmtree(path)
        except Exception as e:
            logger.warning(f"Failed to delete {path}: {e}")

    def _wait_and_delete():
        time.sleep(delay)
        _delete()

    thread_pool_executor.submit(_wait_and_delete)
    atexit.register(_delete)

def create_user_temp_dir():
    session_id = str(uuid.uuid4())[:8]
    temp_dir = os.path.join("temp_local", f"session_{session_id}")
    os.makedirs(temp_dir, exist_ok=True)
    delete_later(temp_dir, delay=600)
    return temp_dir

def load_video_to_tensor(video_path: str) -> torch.Tensor:
    cap = cv2.VideoCapture(video_path)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame)
    cap.release()

    frames = np.array(frames)
    video_tensor = torch.tensor(frames)
    video_tensor = video_tensor.permute(0, 3, 1, 2).float() / 255.0
    video_tensor = video_tensor.unsqueeze(0).permute(0, 2, 1, 3, 4)
    return video_tensor