import gradio as gr
import cv2
import json
import os
import numpy as np
from moviepy.editor import ImageSequenceClip

def process_video(video_path, json_path, threshold):
    # 读取视频
    if os.path.exists("input_video.mp4"):
        video_path = "input_video.mp4"

    cap = cv2.VideoCapture(video_path)
    input_fps = cap.get(cv2.CAP_PROP_FPS)  # 获取输入视频的帧率
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 获取总帧数

    # 读取所有帧
    all_frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # 转换为RGB
        all_frames.append(frame_rgb)
    cap.release()

    # 读取JSON文件
    with open(json_path, 'r') as f:
        mask_area_ratios = json.load(f)

    # 获取mask_area_ratios中的帧索引和对应的ratio
    frame_indices = [int(key.split("_")[1]) for key in mask_area_ratios.keys()]
    ratios = [mask_area_ratios[key] for key in mask_area_ratios.keys()]

    # 创建一个长度为 total_frames 的数组来存储对应的 ratio
    all_ratios = [None] * len(all_frames)

    # 将 mask_area_ratios 的帧索引和对应的 ratio 映射到 all_ratios
    for frame_index, ratio in zip(frame_indices, ratios):
        scaled_index = int(frame_index * (len(all_frames) - 1) / max(frame_indices))
        all_ratios[scaled_index] = ratio

    # 填充缺失的 ratio 值，使用最邻近值
    for i in range(len(all_ratios)):
        if all_ratios[i] is None:
            # 找到最近的非空值
            left_index = i - 1
            right_index = i + 1
            while left_index >= 0 and all_ratios[left_index] is None:
                left_index -= 1
            while right_index < len(all_ratios) and all_ratios[right_index] is None:
                right_index += 1

            if left_index >= 0 and right_index < len(all_ratios):
                # 使用最近的左边或右边的值
                all_ratios[i] = all_ratios[left_index] if all_ratios[left_index] is not None else all_ratios[right_index]
            elif left_index >= 0:
                all_ratios[i] = all_ratios[left_index]
            elif right_index < len(all_ratios):
                all_ratios[i] = all_ratios[right_index]

    # 筛选符合条件的帧
    frames = []
    for frame_idx, ratio in enumerate(all_ratios):
        if ratio >= threshold:
            frames.append(all_frames[frame_idx])

    # 保存符合条件的帧为视频
    if frames:
        output_path = "output_seg_video.mp4"
        clip = ImageSequenceClip(frames, fps=input_fps)  # 使用输入视频的帧率
        clip.write_videofile(output_path, codec="libx264")
        return output_path
    else:
        return "No frames meet the threshold condition."

# Gradio界面
iface = gr.Interface(
    fn=process_video,
    inputs=[
        gr.Video(label="Input Video", value="input_video.mp4"),
        gr.Textbox(label="JSON File Path", value="mask_area_ratios.json"),
        gr.Number(label="Threshold", value=0.25)
    ],
    outputs=gr.Video(label="Output Video"),
    title="Video Frame Filter",
    description="Filter video frames based on mask area ratios and save the result as a new video."
)

iface.launch(share=True)