Update app.py

app.py

@@ -4,6 +4,9 @@ import numpy as np
 import os
 import zipfile
 from PIL import Image
+import tempfile
+import shutil
+from tqdm import tqdm
 
 def get_video_duration(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -13,7 +16,7 @@ def get_video_duration(video_path):
     cap.release()
     return duration
 
-def extract_frames(video_path, output_dir, start_time, end_time):
+def extract_frames(video_path, output_dir, start_time, end_time, max_frames=10000, progress=gr.Progress()):
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS)
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -24,41 +27,46 @@ def extract_frames(video_path, output_dir, start_time, end_time):
     frames = []
     cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
     
-    for i in range(start_frame, min(end_frame, total_frames)):
+    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="프레임 추출 중"):
+        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
         ret, frame = cap.read()
         if not ret:
             break
         frames.append(frame)
+        if len(frames) >= max_frames:
+            break
     
     cap.release()
     return frames
 
-def create_grid(frames, grid_size=(10, 9)):
+def create_grid(frames, grid_size=(10, 10)):
     rows, cols = grid_size
+    n_frames = len(frames)
+    frames = frames[:rows*cols]  # Limit to grid size
+    
     frame_height, frame_width = frames[0].shape[:2]
     grid = np.zeros((frame_height * rows, frame_width * cols, 3), dtype=np.uint8)
     
     for i, frame in enumerate(frames):
-        if i >= rows * cols:
-            break
         row = i // cols
         col = i % cols
         grid[row*frame_height:(row+1)*frame_height, col*frame_width:(col+1)*frame_width] = frame
     
     return grid
 
-def save_frames(frames, output_dir):
+def save_frames(frames, output_dir, progress=gr.Progress()):
     os.makedirs(output_dir, exist_ok=True)
-    for i, frame in enumerate(frames):
+    for i, frame in progress.tqdm(enumerate(frames), desc="프레임 저장 중", total=len(frames)):
         cv2.imwrite(os.path.join(output_dir, f"frame_{i:04d}.jpg"), frame)
 
-def create_zip(output_dir):
+def create_zip(output_dir, progress=gr.Progress()):
     zip_path = os.path.join(output_dir, "frames.zip")
     with zipfile.ZipFile(zip_path, 'w') as zipf:
-        for root, _, files in os.walk(output_dir):
-            for file in files:
-                if file.endswith('.jpg'):
-                    zipf.write(os.path.join(root, file), file)
+        files = [f for f in os.listdir(output_dir) if f.endswith('.jpg')]
+        for file in progress.tqdm(files, desc="ZIP 파일 생성 중"):
+            zipf.write(os.path.join(output_dir, file), file)
     return zip_path
 
 def update_slider(video):
@@ -67,9 +75,7 @@ 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):
+def process_video(video, start_time, end_time, progress=gr.Progress()):
     if video is None:
         return None, None, "비디오 파일을 업로드해주세요."
     
@@ -78,16 +84,20 @@ def process_video(video, start_time, end_time):
         if file_size > 2 * 1024 * 1024 * 1024:  # 2GB in bytes
             return None, None, "파일 크기가 2GB를 초과합니다. 더 작은 파일을 업로드해주세요."
         
-        output_dir = "output_frames"
-        frames = extract_frames(video.name, output_dir, start_time, end_time)
-        if not frames:
-            return None, None, "프레임 추출에 실패했습니다. 비디오 파일을 확인해주세요."
-        
-        save_frames(frames, output_dir)
-        grid = create_grid(frames)
-        zip_path = create_zip(output_dir)
+        with tempfile.TemporaryDirectory() as temp_dir:
+            frames = extract_frames(video.name, temp_dir, start_time, end_time, max_frames=10000, progress=progress)
+            if not frames:
+                return None, None, "프레임 추출에 실패했습니다. 비디오 파일을 확인해주세요."
+            
+            save_frames(frames, temp_dir, progress=progress)
+            grid = create_grid(frames)
+            zip_path = create_zip(temp_dir, progress=progress)
+            
+            # 결과 파일을 임시 디렉토리 밖으로 복사
+            final_zip_path = "frames.zip"
+            shutil.copy(zip_path, final_zip_path)
         
-        return Image.fromarray(cv2.cvtColor(grid, cv2.COLOR_BGR2RGB)), zip_path, "프레임 추출이 완료되었습니다."
+        return Image.fromarray(cv2.cvtColor(grid, cv2.COLOR_BGR2RGB)), final_zip_path, f"프레임 추출이 완료되었습니다. 총 {len(frames)}개의 프레임이 추출되었습니다."
     except Exception as e:
         return None, None, f"오류가 발생했습니다: {str(e)}"
 
@@ -121,7 +131,7 @@ document.addEventListener('DOMContentLoaded', () => {
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, js=js) as iface:
     gr.Markdown("# 비디오 to 이미지(프레임)")
     gr.Markdown("원하는 구간의 프레임을 추출하여 그리드로 표시하고, 개별 이미지를 ZIP 파일로 다운로드할 수 있습니다.")
-    gr.Markdown("최대 2GB 크기의 비디오 파일을 업로드할 수 있습니다.")
+    gr.Markdown("최대 2GB 크기의 비디오 파일을 업로드할 수 있으며, 최대 10000개의 프레임을 추출합니다.")
     
     with gr.Row():
         video_input = gr.Video(label="Upload and Play Video")
@@ -146,4 +156,4 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, js=js) as iface:
     process_button.click(process_video, inputs=[video_input, start_slider, end_slider], outputs=[image_output, file_output, message_output])
 
 if __name__ == "__main__":
-    iface.launch(share=True)  # server_timeout 매개변수 제거
+    iface.launch(share=True)