Update app.py

app.py

@@ -8,6 +8,9 @@ import tempfile
 import shutil
 from tqdm import tqdm
 
+# GPU 사용 가능 여부 확인
+USE_GPU = cv2.cuda.getCudaEnabledDeviceCount() > 0
+
 def get_video_duration(video_path):
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS)
@@ -17,6 +20,12 @@ def get_video_duration(video_path):
     return duration
 
 def extract_frames(video_path, output_dir, start_time, end_time, max_frames=10000, progress=gr.Progress()):
+    if USE_GPU:
+        return extract_frames_gpu(video_path, output_dir, start_time, end_time, max_frames, progress)
+    else:
+        return extract_frames_cpu(video_path, output_dir, start_time, end_time, max_frames, progress)
+
+def extract_frames_cpu(video_path, output_dir, start_time, end_time, max_frames, progress):
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS)
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -29,7 +38,7 @@ def extract_frames(video_path, output_dir, start_time, end_time, max_frames=1000
     
     frame_interval = max(1, (end_frame - start_frame) // max_frames)
     
-    for i in progress.tqdm(range(start_frame, min(end_frame, total_frames), frame_interval), desc="프레임 추출 중"):
+    for i in progress.tqdm(range(start_frame, min(end_frame, total_frames), frame_interval), desc="프레임 추출 중 (CPU)"):
         cap.set(cv2.CAP_PROP_POS_FRAMES, i)
         ret, frame = cap.read()
         if not ret:
@@ -41,6 +50,34 @@ def extract_frames(video_path, output_dir, start_time, end_time, max_frames=1000
     cap.release()
     return frames
 
+def extract_frames_gpu(video_path, output_dir, start_time, end_time, max_frames, progress):
+    stream = cv2.cuda_Stream()
+    cap = cv2.cudacodec.createVideoReader(video_path)
+    
+    _, first_frame = cap.nextFrame()
+    fps = cap.get(cv2.cudacodec.ColorFormat.FPS)
+    total_frames = int(cap.get(cv2.cudacodec.ColorFormat.NUMBER_OF_FRAMES))
+    
+    start_frame = int(start_time * fps)
+    end_frame = int(end_time * fps) if end_time else total_frames
+    
+    frames = []
+    cap.set(cv2.cudacodec.ColorFormat.POS_FRAMES, start_frame)
+    
+    frame_interval = max(1, (end_frame - start_frame) // max_frames)
+    
+    for i in progress.tqdm(range(start_frame, min(end_frame, total_frames), frame_interval), desc="프레임 추출 중 (GPU)"):
+        cap.set(cv2.cudacodec.ColorFormat.POS_FRAMES, i)
+        ret, gpu_frame = cap.nextFrame()
+        if not ret:
+            break
+        cpu_frame = gpu_frame.download()
+        frames.append(cpu_frame)
+        if len(frames) >= max_frames:
+            break
+    
+    return frames
+
 def create_grid(frames, grid_size=(10, 10)):
     rows, cols = grid_size
     n_frames = len(frames)
@@ -75,7 +112,11 @@ def update_slider(video):
     duration = get_video_duration(video.name)
     return gr.Slider(minimum=0, maximum=duration, value=0, step=0.1, label="CUT 시작 시간 (초)"), gr.Slider(minimum=0, maximum=duration, value=duration, step=0.1, label="CUT 종료 시간 (초)")
 
-def process_video(video, start_time, end_time, progress=gr.Progress()):
+
+def process_video(video, start_time, end_time, use_gpu, progress=gr.Progress()):
+    global USE_GPU
+    USE_GPU = use_gpu and cv2.cuda.getCudaEnabledDeviceCount() > 0
+    
     if video is None:
         return None, None, "비디오 파일을 업로드해주세요."
     
@@ -97,37 +138,12 @@ def process_video(video, start_time, end_time, progress=gr.Progress()):
             final_zip_path = "frames.zip"
             shutil.copy(zip_path, final_zip_path)
         
-        return Image.fromarray(cv2.cvtColor(grid, cv2.COLOR_BGR2RGB)), final_zip_path, f"프레임 추출이 완료되었습니다. 총 {len(frames)}개의 프레임이 추출되었습니다."
+        gpu_status = "GPU" if USE_GPU else "CPU"
+        return Image.fromarray(cv2.cvtColor(grid, cv2.COLOR_BGR2RGB)), final_zip_path, f"프레임 추출이 완료되었습니다 ({gpu_status} 사용). 총 {len(frames)}개의 프레임이 추출되었습니다."
     except Exception as e:
         return None, None, f"오류가 발생했습니다: {str(e)}"
 
-css = """
-footer {
-    visibility: hidden;
-}
-#video-time {
-    font-size: 1.2em;
-    font-weight: bold;
-    margin-top: 10px;
-}
-"""
-
-js = """
-function updateVideoTime(videoElement) {
-    const timeDisplay = document.getElementById('video-time');
-    videoElement.addEventListener('timeupdate', () => {
-        const currentTime = videoElement.currentTime.toFixed(2);
-        timeDisplay.textContent = `현재 재생 시간: ${currentTime}초`;
-    });
-}
-document.addEventListener('DOMContentLoaded', () => {
-    const videoElement = document.querySelector('video');
-    if (videoElement) {
-        updateVideoTime(videoElement);
-    }
-});
-"""
-
+# Gradio 인터페이스 설정
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, js=js) as iface:
     gr.Markdown("# 비디오 to 이미지(프레임)")
     gr.Markdown("원하는 구간의 프레임을 추출하여 그리드로 표시하고, 개별 이미지를 ZIP ���일로 다운로드할 수 있습니다.")
@@ -143,6 +159,9 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, js=js) as iface:
         start_slider = gr.Slider(minimum=0, maximum=10000, value=0, step=1, label="CUT 시작 시간 (초)")
         end_slider = gr.Slider(minimum=0, maximum=10000, value=10000, step=1, label="CUT 종료 시간 (초)")
     
+    with gr.Row():
+        use_gpu_checkbox = gr.Checkbox(label="GPU 사용 (가능한 경우)", value=USE_GPU)
+    
     with gr.Row():
         process_button = gr.Button("프레임 추출")
     
@@ -153,7 +172,7 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, js=js) as iface:
     message_output = gr.Textbox(label="메시지")
     
     video_input.change(update_slider, inputs=[video_input], outputs=[start_slider, end_slider])
-    process_button.click(process_video, inputs=[video_input, start_slider, end_slider], outputs=[image_output, file_output, message_output])
+    process_button.click(process_video, inputs=[video_input, start_slider, end_slider, use_gpu_checkbox], outputs=[image_output, file_output, message_output])
 
 if __name__ == "__main__":
     iface.launch(share=True)