Update app.py

app.py

@@ -174,8 +174,17 @@ def remove_bg(fl, count):
     fr_diff = cv2.dilate(cv2.erode(fr_diff, element), element)
 
     if count % 2: # odd: is photo without the flash
-        fr = cv2.imread(fl_).astype(np.uint8)
-        fr = fr_diff - cv2.cvtColor(fr, cv2.COLOR_BGR2GRAY)
+        fr_mask = cv2.imread(fl_).astype(np.uint8)
+        fr_mask = fr_diff - cv2.cvtColor(fr_mask, cv2.COLOR_BGR2GRAY)
+        cv2.imwrite(fl_, fr_mask)
+
+        m = cv2.inRange(fr, np.array([240,240,240]), np.array([255,255,255]))
+        fr[m>0] = (239,239,239)
+        m = cv2.inRange(fr, np.array([0,0,0]), np.array([15,15,15]))
+        fr[m>0] = (16,16,16)
+        fr[fr_mask>0] = (fr[fr_mask>0] / 17).astype(np.uint8)
+        fr[fr_diff==0] = (255,255,255)
+        
         cv2.imwrite(fl, fr)
         return fl
     else: # even: with the flash
@@ -290,7 +299,9 @@ with gr.Blocks() as demo:
     
     gr.Examples(
         examples=[[
-            ["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"], True
+            ["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"], False
+        ], [
+            ["./examples/0_flash.jpg", "./examples/1_noflash.jpg", "./examples/2_flash.jpg", "./examples/3_noflash.jpg"], True
         ]],
         fn=loadf,
         inputs=[files_orig, r_bg],