import cv2
import numpy as np
import os
import gradio as gr

# Функция для увеличения FPS
def increase_fps(input_video, target_fps=60):
    input_path = "input.mp4"
    output_path = "output.mp4"
    
    # Сохранение входного видео
    with open(input_path, "wb") as f:
        f.write(input_video.read())

    cap = cv2.VideoCapture(input_path)
    input_fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    codec = cv2.VideoWriter_fourcc(*"mp4v")

    # Проверка: если входное FPS >= 60, просто копируем файл
    if input_fps >= target_fps:
        os.rename(input_path, output_path)
        with open(output_path, "rb") as f:
            return f.read()

    # Расчет коэффициента интерполяции
    ratio = target_fps / input_fps
    output = cv2.VideoWriter(output_path, codec, target_fps, (width, height))
    _, prev_frame = cap.read()
    prev_frame_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)

    for _ in range(frame_count - 1):
        ret, next_frame = cap.read()
        if not ret:
            break
        next_frame_gray = cv2.cvtColor(next_frame, cv2.COLOR_BGR2GRAY)

        # Оптический поток
        flow = cv2.calcOpticalFlowFarneback(
            prev_frame, next_frame_gray, None, 0.5, 3, 15, 3, 5, 1.2, 0
        )

        # Генерация промежуточных кадров
        for i in range(1, int(ratio)):
            alpha = i / ratio
            intermediate_frame = cv2.addWeighted(prev_frame, 1 - alpha, next_frame_gray, alpha, 0)
            intermediate_bgr = cv2.cvtColor(intermediate_frame, cv2.COLOR_GRAY2BGR)
            output.write(intermediate_bgr)

        output.write(next_frame)
        prev_frame = next_frame_gray

    cap.release()
    output.release()

    # Чтение результата
    with open(output_path, "rb") as f:
        output_video = f.read()

    # Очистка временных файлов
    os.remove(input_path)
    os.remove(output_path)

    return output_video

# Gradio интерфейс
def process_video(video):
    return increase_fps(video)

interface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Загрузите видео для обработки"),
    outputs=gr.Video(label="Видео с увеличенным FPS"),
    title="Увеличение FPS до 60",
    description=(
        "Этот инструмент увеличивает частоту кадров видео до 60 FPS с использованием "
        "оптического потока. Загрузите видео с низким FPS, чтобы получить плавный результат."
    ),
    examples=[]
)

# Запуск интерфейса
if __name__ == "__main__":
    interface.launch()