Update app.py

app.py

@@ -14,6 +14,7 @@ import gradio as gr
 
 from huggingface_hub import snapshot_download
 from image_tools.sizes import resize_and_crop
+from pymatting import cutout
 
 
 model = snapshot_download(repo_id="akhaliq/frame-interpolation-film-style")
@@ -200,9 +201,25 @@ def remove_bg(fl, count, mh, ms, md, lm):
         fr_fg[fr_fg>0] = 3 #probable fg
         mask, bgdModel, fgdModel = cv2.grabCut(fr, fr_fg, None,None,None,31, cv2.GC_INIT_WITH_MASK)
         mask = np.where((mask==2)|(mask==0),0,1).astype('uint8')
-        fr[mask==0] = (255,255,255)
+        #fr[mask==0] = (255,255,255)
         
         cv2.imwrite(fl, fr)
+
+        b = 15
+        element = cv2.getStructuringElement(cv2.MORPH_RECT, (2 * b + 1, 2 * b + 1), (b, b))
+        mask_e = cv2.erode(mask, element) * 255
+        mask_d = cv2.dilate(mask, element) * 127
+        mask_d[mask_e>0] = 255
+        cv2.imwrite(f"{str(count)}_trimask.png", mask_d.astype(np.uint8))
+
+        cutout(
+   # input image path
+   fl,
+   # input trimap path
+   f"{str(count)}_trimask.png",
+   # output cutout path
+   fl)
+        
         return fl
     else: # even: with the flash
         fl_ = fl.split(".")[0] + "_.png"