Update app.py

app.py

@@ -4,27 +4,26 @@ import cv2
 import numpy as np
 import time
 import random
-
 from PIL import Image
 from transparent_background import Remover
 
 @spaces.GPU()
-def doo(video, mode):
-    if(mode == 'Fast'):
+def doo(video, mode, progress=gr.Progress()):
+    if mode == 'Fast':
         remover = Remover(mode='fast')
     else:
         remover = Remover()
+
     cap = cv2.VideoCapture(video)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # Get total frames
     writer = None
     tmpname = random.randint(111111111, 999999999)
     processed_frames = 0
     start_time = time.time()
-    
+
     while cap.isOpened():
         ret, frame = cap.read()
-    
+
         if ret is False:
             break
 
@@ -33,31 +32,33 @@ def doo(video, mode):
         #    cap.release()
         #    writer.release()
         #    return str(tmpname) + '.mp4'
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         img = Image.fromarray(frame).convert('RGB')
-    
+
         if writer is None:
-            writer = cv2.VideoWriter(str(tmpname) + '.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, img.size)
+            writer = cv2.VideoWriter(str(tmpname) + '.mp4', cv2.VideoWriter_fourcc(*'mp4v'), cap.get(cv2.CAP_PROP_FPS), img.size)
 
         processed_frames += 1
-        print(f"Processing: {processed_frames}")
+        print(f"Processing frame {processed_frames}")
+        progress(processed_frames / total_frames, desc=f"Processing frame {processed_frames}/{total_frames}")
         out = remover.process(img, type='green')
         writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
-    
+
     cap.release()
     writer.release()
     return str(tmpname) + '.mp4'
 
-title = "🎞️Video Background Removal tool🎥"
-
-description = r"""Please note that if your video file is long (or having high amount of frames) the result may be shorter than input video because of the GPU timeout.<br>
-In this case consider trying Fast mode."""
+title = "🎞️ Video Background Removal Tool 🎥"
+description = """Please note that if your video file is long (or has a high number of frames), the result may be shorter than the input video due to GPU timeout. In this case, consider trying Fast mode."""
 
-examples = [['./input.mp4'],]
+examples = [['./input.mp4']]
 
 iface = gr.Interface(
-    fn=doo, 
-    inputs=["video", gr.components.Radio(['Normal', 'Fast'], label='Select mode', value='Normal', info= 'Normal is more accurate, but takes longer. | Fast has lower accuracy so the process will be faster.')], 
-    outputs="video", examples=examples, title=title, description=description
+    fn=doo,
+    inputs=["video", gr.components.Radio(['Normal', 'Fast'], label='Select mode', value='Normal', info='Normal is more accurate, but takes longer. | Fast has lower accuracy so the process will be faster.')],
+    outputs="video",
+    examples=examples,
+    title=title,
+    description=description
 )
-iface.launch()
+iface.launch()