Update app.py

app.py

@@ -12,7 +12,10 @@ torch.jit.script = lambda f: f
 from transparent_background import Remover
 
 @spaces.GPU()
-def doo(video, mode, progress=gr.Progress()):
+def doo(video, color, mode, progress=gr.Progress()):
+    color = color.lstrip('#') 
+    rgb = tuple(int(color[i:i+2], 16) for i in (0, 2, 4))
+    color = str(list(rgb))
     if mode == 'Fast':
         remover = Remover(mode='fast')
     else:
@@ -45,7 +48,7 @@ def doo(video, mode, progress=gr.Progress()):
         processed_frames += 1
         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')
+        out = remover.process(img, type=color)
         writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
 
     cap.release()
@@ -59,7 +62,7 @@ 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.')],
+    inputs=["video", gr.ColorPicker(label="Background color", value="#00FF00"), 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,