Update app.py

app.py

@@ -210,7 +210,7 @@ def remove_bg(fl, s, l, v):
     frame[m>0] = (239,239,239)
     m = cv2.inRange(frame, np.array([0,0,0]), np.array([15,15,15]))
     frame[m>0] = (16,16,16)
-    frame[frame_==0] = (frame / 17).astype(np.uint8)
+    frame[frame_==0] = (frame[frame_==0] / 17).astype(np.uint8)
     frame[frame_==255] = (255,255,255)
 
     cv2.imwrite(fl, frame)
@@ -259,7 +259,7 @@ def sortFiles(e):
     e = e.split('/')
     return e[len(e)-1]
 
-def loadf(f, s, l, v):
+def loadf(f, s, l, v, r_bg):
     if f != None and f[0] != None:
         f.sort(key=sortFiles)
         fnew = []
@@ -269,7 +269,9 @@ def loadf(f, s, l, v):
             if ftype[len(ftype)-1].split('.')[1] == 'mp4':
                 fl = sharpest(fl, i)
 
-            fl = remove_bg(fl, s, l, v)
+            if r_bg == True:
+                fl = remove_bg(fl, s, l, v)
+                
             fnew.append(fl)
 
         return fnew, fnew
@@ -313,7 +315,8 @@ with gr.Blocks() as demo:
                         max_s = gr.Slider(minimum=0, maximum=255, step=1, value=32, label="Saturation")
                         max_l = gr.Slider(minimum=0, maximum=255, step=1, value=64, label="Lightness")
                         max_v = gr.Slider(minimum=0, maximum=255, step=1, value=16, label="Variance")
-                files_input.upload(fn=loadf, inputs=[files_input, max_s, max_l, max_v], outputs=[files_input, gallery_input])
+                    rbg = gr.Checkbox(label="Remove background", value=True)
+                files_input.upload(fn=loadf, inputs=[files_input, max_s, max_l, max_v, rbg], outputs=[files_input, gallery_input])
                 with gr.Row():
                     interpolation_slider = gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Interpolation Steps: ")
                     interpolation = gr.Number(value=1, show_label=False, interactive=False)
@@ -331,10 +334,10 @@ with gr.Blocks() as demo:
     gr.Examples(
         examples=[[
             ["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"], 
-            32, 64, 16
+            32, 64, 16, True
         ]],
         fn=loadf,
-        inputs=[files_input, max_s, max_l, max_v],
+        inputs=[files_input, max_s, max_l, max_v, rbg],
         outputs=[files_input, gallery_input],
         cache_examples=True
     )