Update app.py

app.py

@@ -145,7 +145,7 @@ def logscale(linear):
 def linscale(linear):
     return int(math.log2(linear))
 
-def remove_bg(fl, count, mh, ms, md):
+def remove_bg(fl, count, mh, ms, md, lm):
     global fl_
     fr = cv2.imread(fl).astype(np.uint8)
     
@@ -164,6 +164,14 @@ def remove_bg(fl, count, mh, ms, md):
     fr_diff = cv2.cvtColor(fr_diff, cv2.COLOR_BGR2GRAY)
     
     #md = 12
+    if lm == "median":
+        mh = np.median(hsv[:,:,0])
+        ms = np.median(hsv[:,:,1])
+        md = np.median(hsv[:,:,2])
+    elif lm == "average":
+        mh = np.average(hsv[:,:,0])
+        ms = np.average(hsv[:,:,1])
+        md = np.average(hsv[:,:,2])
     bg = cv2.inRange(hsv, np.array([0,0,0]), np.array([mh,ms,md]))
     fr_diff[bg>0] = 0
     fr_diff[bg==0] = 255
@@ -234,7 +242,7 @@ def sortFiles(e):
     e = e.split('/')
     return e[len(e)-1]
 
-def loadf(f, r_bg, mh, ms, md):
+def loadf(f, r_bg, mh, ms, md, lm):
     if f != None and f[0] != None:
         f.sort(key=sortFiles)
         fnew = []
@@ -245,7 +253,7 @@ def loadf(f, r_bg, mh, ms, md):
                 fl = sharpest(fl, i)
 
             if r_bg == True:
-                fl = remove_bg(fl, i, mh, ms, md)
+                fl = remove_bg(fl, i, mh, ms, md, lm)
                 if i % 2: # odd: is photo without the flash
                     fnew.append(fl)
             else:
@@ -288,12 +296,14 @@ with gr.Blocks() as demo:
                     files_orig = gr.File(file_count="multiple", file_types=['image', '.mp4'])
                     files_input = gr.File(file_count="multiple", visible=False)
                 gallery_input = gr.Gallery(label="Slideshow", preview=True, columns=8192, interactive=False)
-                with gr.Accordion(label="Max differences for background", open=False):
-                    mh = gr.Slider(minimum=0, maximum=180, step=1, value=180, label="Hue")
-                    ms = gr.Slider(minimum=0, maximum=255, step=1, value=255, label="Saturation")
-                    md = gr.Slider(minimum=0, maximum=255, step=1, value=12, label="Lightness")
-                r_bg = gr.Checkbox(label="Remove background", value=True)
-                files_orig.upload(fn=loadf, inputs=[files_orig, r_bg, mh, ms, md], outputs=[files_input, gallery_input])
+                with gr.Group(label="Background removal settings"):
+                    r_bg = gr.Checkbox(label="Remove background", value=True)
+                    with gr.Accordion(label="Max differences for background", open=False):
+                        mh = gr.Slider(minimum=0, maximum=180, step=1, value=180, label="Hue")
+                        ms = gr.Slider(minimum=0, maximum=255, step=1, value=255, label="Saturation")
+                        md = gr.Slider(minimum=0, maximum=255, step=1, value=12, label="Lightness")
+                        lm = gr.Radio(label="Use max diffs from", choices=["average", "median", "slider"], value="slider")
+                files_orig.upload(fn=loadf, inputs=[files_orig, r_bg, mh, ms, md, lm], outputs=[files_input, gallery_input])
                 
                 with gr.Row():
                     interpolation_slider = gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Interpolation Steps: ")
@@ -311,12 +321,12 @@ with gr.Blocks() as demo:
     
     gr.Examples(
         examples=[[
-            ["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"], False, 0, 0, 0
+            ["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"], False, 0, 0, 0, "slider"
         ], [
-            ["./examples/0_flash.jpg", "./examples/1_noflash.jpg", "./examples/2_flash.jpg", "./examples/3_noflash.jpg"], True, 180, 255, 12
+            ["./examples/0_flash.jpg", "./examples/1_noflash.jpg", "./examples/2_flash.jpg", "./examples/3_noflash.jpg"], True, 180, 255, 12, "slider"
         ]],
         fn=loadf,
-        inputs=[files_orig, r_bg, mh, ms, md],
+        inputs=[files_orig, r_bg, mh, ms, md, lm],
         outputs=[files_input, gallery_input],
         cache_examples=True
     )