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):
+def remove_bg(fl, count, md):
     global fl_
     fr = cv2.imread(fl).astype(np.uint8)
     
@@ -159,11 +159,14 @@ def remove_bg(fl, count):
         fr_bg = cv2.medianBlur(fr_bg, 255)
         
     fr_diff = cv2.convertScaleAbs(fr.astype(np.int16)-fr_bg.astype(np.int16)).astype(np.uint8)
+    hsv = cv2.cvtColor(fr_diff, cv2.COLOR_BGR2HSV) # range: 180, 255, 255
+    
     fr_diff = cv2.cvtColor(fr_diff, cv2.COLOR_BGR2GRAY)
     
-    md = 12
-    fr_diff[fr_diff>=md] = 255
-    fr_diff[fr_diff<md] = 0
+    #md = 12
+    fg = np.bitwise_and(hsv[:,:,2]>=md, hsv[:,:,1]>=md, hsv[:,:,0]>=md)
+    fr_diff[fg>0] = 255
+    fr_diff[fg==0] = 0
     
     cv2.rectangle(fr_diff,(0,0),(fr_diff.shape[1]-1,fr_diff.shape[0]-1),(255,255,255),1)
     mask = cv2.floodFill(fr_diff, None, (0, 0), 255, 0, 0, (4 | cv2.FLOODFILL_FIXED_RANGE))[2] #(4 | cv.FLOODFILL_FIXED_RANGE | cv.FLOODFILL_MASK_ONLY | 255 << 8)
@@ -231,7 +234,7 @@ def sortFiles(e):
     e = e.split('/')
     return e[len(e)-1]
 
-def loadf(f, r_bg):
+def loadf(f, r_bg, md):
     if f != None and f[0] != None:
         f.sort(key=sortFiles)
         fnew = []
@@ -242,7 +245,7 @@ def loadf(f, r_bg):
                 fl = sharpest(fl, i)
 
             if r_bg == True:
-                fl = remove_bg(fl, i)
+                fl = remove_bg(fl, i, md)
                 if i % 2: # odd: is photo without the flash
                     fnew.append(fl)
             else:
@@ -286,7 +289,8 @@ with gr.Blocks() as demo:
                     files_input = gr.File(file_count="multiple", visible=False)
                 gallery_input = gr.Gallery(label="Slideshow", preview=True, columns=8192, interactive=False)
                 r_bg = gr.Checkbox(label="Remove background", value=True)
-                files_orig.upload(fn=loadf, inputs=[files_orig, r_bg], outputs=[files_input, gallery_input])
+                md = gr.Slider(minimum=0, maximum=255, step=1, value=12, label="Min color diff for foreground")
+                files_orig.upload(fn=loadf, inputs=[files_orig, r_bg, md], outputs=[files_input, gallery_input])
                 
                 with gr.Row():
                     interpolation_slider = gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Interpolation Steps: ")
@@ -304,12 +308,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
+            ["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"], False, 0
         ], [
-            ["./examples/0_flash.jpg", "./examples/1_noflash.jpg", "./examples/2_flash.jpg", "./examples/3_noflash.jpg"], True
+            ["./examples/0_flash.jpg", "./examples/1_noflash.jpg", "./examples/2_flash.jpg", "./examples/3_noflash.jpg"], True, 12
         ]],
         fn=loadf,
-        inputs=[files_orig, r_bg],
+        inputs=[files_orig, r_bg, md],
         outputs=[files_input, gallery_input],
         cache_examples=True
     )