Update app.py

app.py

@@ -630,7 +630,7 @@ def process_bg_image(image, bg, fast_mode=False):
     image = Image.composite(image, background, mask)
     return image
 
-def process_video_frame(frame, bg_type, bg, fast_mode, bg_frame_index, background_frames, color):
+def process_video_frame(frame, bg_type, bg, fast_mode, frame_index, background_frames, color):
     """비디오 프레임 처리"""
     try:
         pil_image = Image.fromarray(frame)
@@ -639,16 +639,29 @@ def process_video_frame(frame, bg_type, bg, fast_mode, bg_frame_index, backgroun
         elif bg_type == "이미지":
             processed_image = process_bg_image(pil_image, bg, fast_mode)
         elif bg_type == "비디오":
-            background_frame = background_frames[bg_frame_index]
-            bg_frame_index += 1
-            background_image = Image.fromarray(background_frame)
-            processed_image = process_bg_image(pil_image, background_image, fast_mode)
+            # 인덱스 범위 확인
+            if background_frames and len(background_frames) > 0:
+                # 프레임 인덱스를 배경 비디오 길이로 나눈 나머지를 사용 (루프 효과)
+                bg_frame_index = frame_index % len(background_frames)
+                background_frame = background_frames[bg_frame_index]
+                background_image = Image.fromarray(background_frame)
+                processed_image = process_bg_image(pil_image, background_image, fast_mode)
+            else:
+                processed_image = pil_image
         else:
             processed_image = pil_image
-        return np.array(processed_image), bg_frame_index
+        
+        # 처리된 이미지가 numpy array로 반환되는지 확인
+        if isinstance(processed_image, Image.Image):
+            return np.array(processed_image)
+        return processed_image
+        
     except Exception as e:
-        print(f"Error processing frame: {e}")
-        return frame, bg_frame_index
+        print(f"Error processing frame {frame_index}: {e}")
+        # 오류 발생 시 원본 프레임 반환
+        if isinstance(frame, np.ndarray):
+            return frame
+        return np.array(pil_image)
 
 @spaces.GPU(duration=300)
 def process_video_bg(vid, bg_type="색상", bg_image=None, bg_video=None, color="#00FF00", 
@@ -672,44 +685,104 @@ def process_video_bg(vid, bg_type="색상", bg_image=None, bg_video=None, color=
         audio = video.audio
         frames = list(video.iter_frames(fps=fps))
         
+        # 프레임 크기 저장
+        if frames:
+            frame_height, frame_width = frames[0].shape[:2]
+        else:
+            yield gr.update(visible=False), gr.update(visible=True), "비디오에 프레임이 없습니다."
+            yield None, None, "비디오에 프레임이 없습니다."
+            return
+        
         processed_frames = []
         yield gr.update(visible=True), gr.update(visible=False), f"처리 시작... 경과 시간: 0초"
 
-        if bg_type == "비디오":
+        # 배경 비디오 처리
+        background_frames = None
+        if bg_type == "비디오" and bg_video:
             background_video = VideoFileClip(bg_video)
-            if background_video.duration < video.duration:
-                if video_handling == "slow_down":
-                    # vfx.speedx 대신 speedx 함수 직접 사용
-                    if speedx is not None:
-                        background_video = speedx(background_video, factor=video.duration / background_video.duration)
-                    else:
-                        # speedx가 없으면 반복으로 대체
-                        background_video = concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
-                else:  # video_handling == "loop"
-                    background_video = concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
+            
+            # 배경 비디오 길이 조정
+            if video_handling == "slow_down" and background_video.duration < video.duration:
+                if speedx is not None:
+                    factor = video.duration / background_video.duration
+                    background_video = speedx(background_video, factor=factor)
+                else:
+                    # speedx가 없으면 반복으로 대체
+                    loops = int(video.duration / background_video.duration) + 1
+                    background_video = concatenate_videoclips([background_video] * loops)
+            elif video_handling == "loop" or background_video.duration < video.duration:
+                # 반복 모드
+                loops = int(video.duration / background_video.duration) + 1
+                background_video = concatenate_videoclips([background_video] * loops)
+            
+            # 배경 프레임 추출
             background_frames = list(background_video.iter_frames(fps=fps))
-        else:
-            background_frames = None
-
-
-        
-        bg_frame_index = 0
+            
+            # 배�� 비디오가 더 길면 잘라냄
+            if len(background_frames) > len(frames):
+                background_frames = background_frames[:len(frames)]
 
+        # 병렬 처리
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
-            futures = [executor.submit(process_video_frame, frames[i], bg_type, bg_image, fast_mode, 
-                                     bg_frame_index + i, background_frames, color) for i in range(len(frames))]
+            futures = []
+            for i in range(len(frames)):
+                future = executor.submit(
+                    process_video_frame, 
+                    frames[i], 
+                    bg_type, 
+                    bg_image, 
+                    fast_mode, 
+                    i,  # 프레임 인덱스 전달
+                    background_frames, 
+                    color
+                )
+                futures.append(future)
+            
+            # 결과 수집
             for i, future in enumerate(futures):
-                result, _ = future.result()
-                processed_frames.append(result)
-                elapsed_time = time.time() - start_time
-                yield result, None, f"프레임 {i+1}/{len(frames)} 처리 중... 경과 시간: {elapsed_time:.2f}초"
+                try:
+                    result = future.result()
+                    # 결과가 올바른 크기인지 확인
+                    if result.shape[:2] != (frame_height, frame_width):
+                        # 크기가 다르면 리사이즈
+                        pil_result = Image.fromarray(result)
+                        pil_result = pil_result.resize((frame_width, frame_height), Image.LANCZOS)
+                        result = np.array(pil_result)
+                    
+                    processed_frames.append(result)
+                    elapsed_time = time.time() - start_time
+                    
+                    # 10프레임마다 상태 업데이트
+                    if i % 10 == 0:
+                        yield result, None, f"프레임 {i+1}/{len(frames)} 처리 중... 경과 시간: {elapsed_time:.2f}초"
+                except Exception as e:
+                    print(f"Error getting result for frame {i}: {e}")
+                    # 오류 발생 시 원본 프레임 사용
+                    processed_frames.append(frames[i])
+        
+        # 모든 프레임이 동일한 크기인지 최종 확인
+        frame_sizes = [frame.shape for frame in processed_frames]
+        if len(set(frame_sizes)) > 1:
+            print(f"Warning: Different frame sizes detected: {set(frame_sizes)}")
+            # 첫 번째 프레임 크기로 모두 통일
+            target_size = processed_frames[0].shape
+            for i in range(len(processed_frames)):
+                if processed_frames[i].shape != target_size:
+                    pil_frame = Image.fromarray(processed_frames[i])
+                    pil_frame = pil_frame.resize((target_size[1], target_size[0]), Image.LANCZOS)
+                    processed_frames[i] = np.array(pil_frame)
         
+        # 비디오 생성
         processed_video = ImageSequenceClip(processed_frames, fps=fps)
-        processed_video = processed_video.with_audio(audio)
         
+        # 오디오 추가
+        if audio:
+            processed_video = processed_video.set_audio(audio)
+        
+        # 저장
         with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
             temp_filepath = temp_file.name
-            processed_video.write_videofile(temp_filepath, codec="libx264")
+            processed_video.write_videofile(temp_filepath, codec="libx264", audio_codec="aac")
         
         elapsed_time = time.time() - start_time
         yield gr.update(visible=False), gr.update(visible=True), f"처리 완료! 경과 시간: {elapsed_time:.2f}초"
@@ -717,6 +790,8 @@ def process_video_bg(vid, bg_type="색상", bg_image=None, bg_video=None, color=
     
     except Exception as e:
         print(f"Error: {e}")
+        import traceback
+        traceback.print_exc()
         elapsed_time = time.time() - start_time
         yield gr.update(visible=False), gr.update(visible=True), f"비디오 처리 오류: {e}. 경과 시간: {elapsed_time:.2f}초"
         yield None, None, f"비디오 처리 오류: {e}. 경과 시간: {elapsed_time:.2f}초"