Update app.py

app.py

@@ -140,21 +140,24 @@ def remove_bg(fl):
     frame = cv2.imread(fl).astype(np.uint8)
     
     b = 5
-    #subtract background (get scene with shadow) and correct hue against light
+    #subtract background (get scene with shadow)
     bg = cv2.medianBlur(frame, 255)
+    
+    frame_ = ((bg.astype(np.int16)-frame.astype(np.int16))+127).astype(np.uint8)
+    frame_ = cv2.bilateralFilter(frame_, b*4, b*8, b*2)
+    frame_ = cv2.medianBlur(frame_, b)
+    
+    element = cv2.getStructuringElement(cv2.MORPH_RECT, (2*b+1, 2*b+1), (b,b))
+    frame_ = cv2.erode(cv2.dilate(frame_, element), element)
+
+    #correct hue against light
     bg_gray = cv2.cvtColor(cv2.cvtColor(bg, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR)
     bg_diff = (bg-bg_gray).astype(np.int16)
-    frame_ = ((bg.astype(np.int16)-frame.astype(np.int16))+127).astype(np.uint8)
     frame = (frame.astype(np.int16)-bg_diff).astype(np.uint8)
-    frame_ = cv2.bilateralFilter(frame_, b*4, b*8, b*2)
     
     #remove regions of low saturation (get scene without shadow)
-    mask = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
-    m = cv2.inRange(mask, np.array([0,0,0]), np.array([180,32,256]))
+    m = cv2.inRange(cv2.cvtColor(frame, cv2.COLOR_RGB2HSV), np.array([0,0,0]), np.array([180,32,256]))
     frame[m>0] = (127,127,127)
-    frame_ = cv2.medianBlur(frame_, b)
-    element = cv2.getStructuringElement(cv2.MORPH_RECT, (2*b+1, 2*b+1), (b,b))
-    frame_ = cv2.erode(cv2.dilate(frame_, element), element)
     frame = cv2.medianBlur(frame, b)
     
     
@@ -164,7 +167,7 @@ def remove_bg(fl):
     # define criteria, number of clusters(K) and apply kmeans()
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 16, 1.0)
     K = 3
-    ret,label,center=cv2.kmeans(m_,K,None,criteria,16,cv2.KMEANS_PP_CENTERS)
+    ret,label,center=cv2.kmeans(m_, K, None, criteria, 16, cv2.KMEANS_PP_CENTERS)
     # Now convert back into uint8, and make original image
     center = np.uint8(center)
     res = center[label.flatten()]
@@ -172,15 +175,18 @@ def remove_bg(fl):
     
     
     #remove shadows at edges
-    cv2.rectangle(frame_,(0,0),(frame_.shape[1]-1,frame_.shape[0]-1),(142,142,142),7)
+    m_ = cv2.inRange(frame_, np.array([128,128,128]), np.array([255,255,255]))
+    frame_[m_>0] = (255,255,255)
+    cv2.rectangle(frame_,(0,0),(frame_.shape[1]-1,frame_.shape[0]-1),(255,255,255),7)
     mask = cv2.floodFill(frame_, None, (0, 0), 255, 0, 0, (4 | cv2.FLOODFILL_FIXED_RANGE))[2] #(4 | cv2.FLOODFILL_FIXED_RANGE | cv2.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]
-    frame_[mask>0] = (125,125,125)
-    frame_[frame_[:,:,0]>=142] = (255,255,255)
+    frame_[mask>0] = (127,127,127)
+    m_ = cv2.inRange(frame_, np.array([0,0,0]), np.array([127,127,127]))
+    frame_[m_>0] = (127,127,127)
+    
 
     cv2.imwrite(fl, frame) #frame_
-
     return fl
 
 def logscale(linear):
@@ -292,7 +298,7 @@ with gr.Blocks() as demo:
     
     gr.Examples(
         examples=[[["./examples/0.png", "./examples/1.png", "./examples/2.png", "./examples/3.png", "./examples/4.png"], 1, 0]],
-        fn=infer,
+        fn=loadf,
         inputs=[files_input, interpolation_slider, fps_output_slider],
         outputs=[video_output, file_output],
         cache_examples=True