import cv2
import numpy as np
import os

# 設(shè)置原始視頻文件夾路徑.
input_folder = "C:/Users/Administrator/Desktop/output"  # 替換為你的輸入文件夾路徑.
output_file = 'C:/Users/Administrator/Desktop/combined_video.mp4'   # 輸出文件名.

# 指定視頻的順序.
video_indices = [1, 2, 3]  # 假設(shè)你的視頻編號(hào)從1到3.

# 收集視頻文件名.
original_videos = []
segmentation_videos = []

for index in video_indices:
    original_videos.append(os.path.join(input_folder, f'resized_test_video{index}.mp4'))
    segmentation_videos.append(os.path.join(input_folder, f'resized_output_video{index}.mp4'))

# 確保視頻數(shù)量匹配.
if len(original_videos) != len(segmentation_videos):
    print("原始視頻和分割視頻數(shù)量不匹配！")
    exit()

# 初始化VideoCapture對(duì)象.
caps = [cv2.VideoCapture(v) for v in original_videos + segmentation_videos]

# 獲取視頻幀的寬高（確保所有視頻具有相同分辨率）.
frame_width = int(caps[0].get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(caps[0].get(cv2.CAP_PROP_FRAME_HEIGHT))

# 設(shè)置輸出視頻的參數(shù).
output_width = frame_width * 3
output_height = frame_height * 2
fps = caps[0].get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')

# 創(chuàng)建 VideoWriter 對(duì)象.
out = cv2.VideoWriter(output_file, fourcc, fps, (output_width, output_height))

while True:
    frames = []

    # 讀取所有視頻的當(dāng)前幀.
    for cap in caps:
        ret, frame = cap.read()
        if not ret:
            frames.append(None)  # 添加空幀以保持一致性.
        else:
            # Resize frame if needed
            frame = cv2.resize(frame, (frame_width, frame_height))
            frames.append(frame)

    # 檢查是否所有視頻都已經(jīng)結(jié)束.
    if all(frame is None for frame in frames):
        break

    # 創(chuàng)建一個(gè)大畫面，只合并有效幀.
    valid_frames = [frame for frame in frames if frame is not None]

    # 檢查是否有足夠的有效幀進(jìn)行合并.
    if len(valid_frames) == 6:
        top_row = np.hstack((valid_frames[0], valid_frames[1], valid_frames[2]))
        bottom_row = np.hstack((valid_frames[3], valid_frames[4], valid_frames[5]))
        combined_frame = np.vstack((top_row, bottom_row))

        # 寫入合成幀到輸出文件.
        out.write(combined_frame)
    else:
        print("當(dāng)前幀無(wú)效，無(wú)法進(jìn)行合并。")

# 釋放資源.
for cap in caps:
    cap.release()
out.release()

print('視頻合并完成！輸出文件:', output_file)

# 實(shí)現(xiàn)效果.
from IPython.display import Video

Video(filename='./輸出效果/combined_video.mp4', width=450, height=330)