api/predict.py

import os
import sys
import time
import tempfile
import uuid

import torch
import numpy as np
from flask import Blueprint, request, jsonify, current_app

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config import config
from utils.video_utils import load_video_frames, normalize_frames, validate_video

predict_bp = Blueprint("predict", __name__)

ALLOWED_EXTENSIONS = {".mp4", ".avi", ".mkv", ".mov", ".webm"}


def _allowed_file(filename: str) -> bool:
    return os.path.splitext(filename.lower())[1] in ALLOWED_EXTENSIONS


@predict_bp.route("/predict", methods=["POST"])
def predict():
    # ── Validate upload ───────────────────────────────────────────────────────
    if "video" not in request.files:
        return jsonify({"error": "No video file provided"}), 400

    file = request.files["video"]
    if file.filename == "":
        return jsonify({"error": "Empty filename"}), 400

    if not _allowed_file(file.filename):
        exts = ", ".join(ALLOWED_EXTENSIONS)
        return jsonify({"error": f"Unsupported format. Allowed: {exts}"}), 415

    # ── Save to temp path ─────────────────────────────────────────────────────
    suffix = os.path.splitext(file.filename)[1]
    tmp_path = os.path.join(tempfile.gettempdir(), f"aivd_{uuid.uuid4().hex}{suffix}")
    try:
        file.save(tmp_path)
    except Exception as e:
        return jsonify({"error": f"Could not save file: {e}"}), 500

    # ── Inference ─────────────────────────────────────────────────────────────
    t_start = time.time()
    try:
        model = current_app.model
        device = current_app.device

        if not validate_video(tmp_path):
            return jsonify({"error": "Invalid or corrupted video file"}), 422

        frames = load_video_frames(
            tmp_path,
            num_frames=config.FRAMES_PER_VIDEO,
            frame_size=(config.FRAME_HEIGHT, config.FRAME_WIDTH),
            frame_skip=config.FRAME_SKIP,
        )
        if frames is None:
            return jsonify({"error": "Could not extract frames from video"}), 422

        frames = normalize_frames(frames)

        frames_tensor = torch.from_numpy(frames).float()
        frames_tensor = frames_tensor.permute(0, 3, 1, 2)   # (T, 3, H, W)
        frames_tensor = frames_tensor.unsqueeze(0)            # (1, T, 3, H, W)
        frames_tensor = frames_tensor.to(device)

        with torch.no_grad():
            output = model(frames_tensor)
            prob = torch.sigmoid(output).item()

        threshold = config.PREDICTION_THRESHOLD
        is_ai = prob > threshold
        verdict = "AI GENERATED" if is_ai else "REAL VIDEO"
        confidence = prob if is_ai else (1.0 - prob)
        elapsed = round(time.time() - t_start, 2)

        return jsonify({
            "verdict": verdict,
            "is_ai": is_ai,
            "probability": round(prob, 6),
            "confidence": round(confidence * 100, 2),
            "processing_time": elapsed,
            "threshold": threshold,
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500

    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass