Update app.py

app.py

@@ -162,21 +162,18 @@ def remove_bg(fl, count, mh, ms, md, lm):
     hsv = cv2.cvtColor(fr_diff, cv2.COLOR_BGR2HSV) # range: 180, 255, 255
     
     fr_diff = cv2.cvtColor(fr_diff, cv2.COLOR_BGR2GRAY)
-
-    fr_diff = cv2.adaptiveThreshold(fr_diff,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,255,0)
     
-    #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
+    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
     
     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)
@@ -199,6 +196,11 @@ def remove_bg(fl, count, mh, ms, md, lm):
         fr[m>0] = (16,16,16)
         fr[fr_shadow>0] = (fr[fr_shadow>0] / 17).astype(np.uint8)
         fr[fr_fg==0] = (255,255,255)
+
+        fr_fg[fr_fg>0] = 3 #probable fg
+        mask, bgdModel, fgdModel = cv2.grabCut(fr, fr_fg, None,None,None,15, cv2.GC_INIT_WITH_MASK)
+        mask = np.where((mask==2)|(mask==0),0,1).astype('uint8')
+        fr[mask==0] = (255,255,255)
         
         cv2.imwrite(fl, fr)
         return fl