Update app.py

app.py

@@ -46,7 +46,8 @@ def doo(video, mode, progress=gr.Progress()):
         if time.time() - start_time >= 20 * 60 - 5:
             print("GPU Timeout is coming")
             cap.release()
-            writer.release()
+            if writer:
+                writer.release()
             return str(tmpname) + '.mp4'
 
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
@@ -62,8 +63,10 @@ def doo(video, mode, progress=gr.Progress()):
         # 배경 제거
         out = remover.process(img, type='green')
         
-        # 마스크 추출 및 후처리
-        mask = np.array(out)[:,:,3]
+        # 마스크 생성 (RGB to grayscale)
+        mask = cv2.cvtColor(np.array(out), cv2.COLOR_RGB2GRAY)
+        
+        # 마스크 후처리
         mask = post_process_mask(mask)
         
         # 시간적 평활화 적용
@@ -71,13 +74,15 @@ def doo(video, mode, progress=gr.Progress()):
         previous_mask = mask
 
         # 마스크 적용 및 색상 보정
-        result = cv2.multiply(frame, cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB).astype(float) / 255.0)
+        mask_3d = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
+        result = cv2.multiply(frame, mask_3d.astype(float) / 255.0)
         result = cv2.addWeighted(result, 1.1, frame, 0, 0)  # 색상 보정
 
         writer.write(cv2.cvtColor(result.astype(np.uint8), cv2.COLOR_RGB2BGR))
 
     cap.release()
-    writer.release()
+    if writer:
+        writer.release()
     return str(tmpname) + '.mp4'
 
 title = "🎞️ Enhanced Video Background Removal Tool 🎥"