Update app.py

app.py

@@ -23,6 +23,9 @@ interpolator = interpolator.Interpolator(model, None)
 ffmpeg_path = util.get_ffmpeg_path()
 mediapy.set_ffmpeg(ffmpeg_path)
 
+fl_ = ""
+fl_mask = ""
+
 
 def do_interpolation(frame1, frame2, interpolation, n):
     print("tween frames: " + str(interpolation))
@@ -142,6 +145,43 @@ def logscale(linear):
 def linscale(linear):
     return int(math.log2(linear))
 
+def remove_bg(fl, i):
+    global fl_
+    fr = cv2.imread(fl)
+    
+    b = 1
+    element = cv2.getStructuringElement(cv2.MORPH_RECT, (2 * b + 1, 2 * b + 1), (b, b))
+    
+    n = int((fr.shape[0]*fr.shape[1]) / (256*256))
+    fr_bg = cv2.medianBlur(fr, 255)
+    
+    for i in range(0, n):
+        fr_bg = cv2.medianBlur(fr_bg, 255)
+        
+    fr_diff = cv2.convertScaleAbs(fr.astype(np.int16)-fr_bg.astype(np.int16)).astype(np.uint8)
+    fr_diff = cv2.cvtColor(fr_diff, cv2.COLOR_BGR2GRAY)
+    
+    md = 12
+    fr_diff[fr_diff>=md] = 255
+    fr_diff[fr_diff<md] = 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)
+    # 255 << 8 tells to fill with the value 255)
+    mask = mask[1:mask.shape[0]-1, 1:mask.shape[1]-1]
+    fr_diff[mask>0] = 0
+    
+    fr_diff = cv2.dilate(cv2.erode(fr_diff, element), element)
+
+    if int(i/2) == i/2: # is photo with the flash
+        fl_ = fl.split(".")[0] + "_.png"
+        cv2.imwrite(fl_, fr_diff)
+    else: # is without
+        fr = cv2.imread(fl_)
+        fr = fr_diff - fr
+        cv2.imwrite(fl_, fr)
+    return fl_
+
 def sharpest(fl, i):
     break_vid = get_frames(fl, "vid_input_frame", "origin", i)
 
@@ -179,7 +219,7 @@ def sortFiles(e):
     e = e.split('/')
     return e[len(e)-1]
 
-def loadf(f):
+def loadf(f, r_bg):
     if f != None and f[0] != None:
         f.sort(key=sortFiles)
         fnew = []
@@ -188,6 +228,9 @@ def loadf(f):
             ftype = fl.split('/')
             if ftype[len(ftype)-1].split('.')[1] == 'mp4':
                 fl = sharpest(fl, i)
+
+            if r_bg == True:
+                fl = remove_bg(fl, i)
                 
             fnew.append(fl)
         return fnew, fnew
@@ -227,7 +270,8 @@ 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)
-                files_orig.upload(fn=loadf, inputs=[files_orig], outputs=[files_input, gallery_input])
+                r_bg = gr.Checkbox(label="Remove background", value=True)
+                files_orig.upload(fn=loadf, inputs=[files_orig, r_bg], outputs=[files_input, gallery_input])
                 
                 with gr.Row():
                     interpolation_slider = gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Interpolation Steps: ")
@@ -245,10 +289,10 @@ with gr.Blocks() as demo:
     
     gr.Examples(
         examples=[[
-            ["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"]
+            ["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"], True
         ]],
         fn=loadf,
-        inputs=[files_orig],
+        inputs=[files_orig, r_bg],
         outputs=[files_input, gallery_input],
         cache_examples=True
     )