Update app.py

app.py

@@ -79,7 +79,7 @@ def predict_depth(image, model):
 #def predict_depth(model, image):
 #    return model(image)["depth"]
 
-def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg=False, maxc=16, maxd=16, lt="slider"):
+def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg=False, maxc=12, maxd=12, maxs=32, maxl=64, maxv=16, lt="slider"):
     if encoder not in ["vitl","vitb","vits","vitg"]:
         encoder = "vits"
 
@@ -232,7 +232,29 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
                     
                     diff_d = np.abs(depth_color.astype(np.int16)-depth_color_bg.astype(np.int16))
                     diff_c = np.abs(raw_frame.astype(np.int16)-raw_frame_bg.astype(np.int16))
-                    
+
+                    #correct hue against light
+                    bg_gray = cv2.cvtColor(cv2.cvtColor(raw_frame_bg, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR)
+                    bg_diff = (raw_frame_bg-bg_gray).astype(np.int16)
+                    frame_c = np.abs(raw_frame.astype(np.int16)-bg_diff).astype(np.uint8)
+
+                    hsv_ = cv2.cvtColor(frame_c, cv2.COLOR_BGR2HSV)
+                    edges = cv2.Laplacian(cv2.cvtColor(raw_frame, cv2.COLOR_BGR2GRAY), cv2.CV_64F)
+                    blur_s = np.zeros_like(edges)
+                    for i in range(2, frame.shape[0]-2):
+                        for j in range(2, frame.shape[1]-2):
+                            d = edges[i-2:i+2, j-2:j+2].var()
+                            blur_s[i,j] = d.astype(np.uint8)
+
+                    print("detail")
+                    print(np.average(blur_s))
+                    print(np.median(blur_s))
+                    print("saturation")
+                    print(np.average(hsv_[:,:,1]))
+                    print(np.median(hsv_[:,:,1]))
+                    print("lightness")
+                    print(np.average(hsv_[:,:,2]))
+                    print(np.median(hsv_[:,:,2]))
                     #print('-most common')
                     #c = Counter(diff_d.flatten())
                     #value, cc = c.most_common()[0]
@@ -252,6 +274,9 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
                         print(md_c)
                         mask_bg_shadow = cv2.inRange(diff_d, np.array([0,0,0]), np.array([md_d,md_d,md_d]))
                         mask_bg_no_shadow = cv2.inRange(diff_c, np.array([0,0,0]), np.array([md_c,md_c,md_c]))
+
+                        m = cv2.inRange(hsv_, np.array([0,0,0]), np.array([180, int(np.median(hsv_[:,:,1])), int(np.median(hsv_[:,:,2]))]))
+                        mask = cv2.inRange(blur_s, 0, int(np.median(blur_s)))
                     elif lt == "average":
                         avg_d = int(np.average(diff_d))
                         avg_c = int(np.average(diff_c))
@@ -260,9 +285,18 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
                         print(avg_c)
                         mask_bg_shadow = cv2.inRange(diff_d, np.array([0,0,0]), np.array([avg_d,avg_d,avg_d]))
                         mask_bg_no_shadow = cv2.inRange(diff_c, np.array([0,0,0]), np.array([avg_c,avg_c,avg_c]))
+
+                        m = cv2.inRange(hsv_, np.array([0,0,0]), np.array([180, int(np.average(hsv_[:,:,1])), int(np.average(hsv_[:,:,2]))]))
+                        mask = cv2.inRange(blur_s, 0, int(np.average(blur_s)))
                     elif lt == "slider":
                         mask_bg_shadow = cv2.inRange(diff_d, np.array([0,0,0]), np.array([maxd,maxd,maxd]))
                         mask_bg_no_shadow = cv2.inRange(diff_c, np.array([0,0,0]), np.array([maxc,maxc,maxc]))
+
+                        m = cv2.inRange(hsv_, np.array([0,0,0]), np.array([180,maxs,maxl]))
+                        mask = cv2.inRange(blur_s, 0, maxv)
+
+                    masks = np.bitwise_and(m, mask)
+                    masks_shadow = np.bitwise_and(mask_bg_shadow, masks)
                     #mask_no_shadow = cv2.bitwise_not(mask_shadow)
 
                     #stereo = cv2.StereoBM.create(numDisparities=16, blockSize=15)
@@ -272,7 +306,7 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
                     raw_frame[m>0] = (239,239,239)
                     m = cv2.inRange(raw_frame, np.array([0,0,0]), np.array([15,15,15]))
                     raw_frame[m>0] = (16,16,16)
-                    raw_frame[mask_bg_shadow>0] = (raw_frame[mask_bg_shadow>0] / 17).astype(np.uint8)
+                    raw_frame[masks_shadow>0] = (raw_frame[masks_shadow>0] / 17).astype(np.uint8)
                     raw_frame[mask_bg_no_shadow>0] = (255,255,255)
                 else:
                     break
@@ -1046,6 +1080,10 @@ with gr.Blocks(css=css, js=js) as demo:
                 with gr.Tab(label="Maximums"):
                     max_c = gr.Slider(minimum=0, maximum=255, step=1, value=12, label="Color diff")
                     max_d = gr.Slider(minimum=0, maximum=255, step=1, value=12, label="Depth diff")
+                with gr.Tab(label="Shadow maximums"):
+                    max_s = gr.Slider(minimum=0, maximum=255, step=1, value=32, label="Saturation")
+                    max_l = gr.Slider(minimum=0, maximum=255, step=1, value=64, label="Lightness")
+                    max_v = gr.Slider(minimum=0, maximum=255, step=1, value=16, label="Detail")
                 lt = gr.Radio(label="Maximum is", choices=["average", "median", "slider"], value="slider")
             processed_video = gr.Video(label="Output Video", format="mp4", interactive=False)
             processed_zip = gr.File(label="Output Archive", interactive=False)
@@ -1383,7 +1421,7 @@ with gr.Blocks(css=css, js=js) as demo:
             render = gr.Button("Render")
             input_json.input(show_json, inputs=[input_json], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
     
-    def on_submit(uploaded_video,model_type,remove_bg,maxc,maxd,lt,coordinates):
+    def on_submit(uploaded_video,model_type,remove_bg,maxc,maxd,maxs,maxl,maxv,lt,coordinates):
         global locations
         locations = []
         avg = [0, 0]
@@ -1417,16 +1455,16 @@ with gr.Blocks(css=css, js=js) as demo:
         print(locations)
             
         # Process the video and get the path of the output video
-        output_video_path = make_video(uploaded_video,encoder=model_type,remove_bg=remove_bg,maxc=maxc,maxd=maxd,lt=lt)
+        output_video_path = make_video(uploaded_video,encoder=model_type,remove_bg=remove_bg,maxc=maxc,maxd=maxd,maxs=maxs,maxl=maxl,maxv=maxv,lt=lt)
 
         return output_video_path + (json.dumps(locations),)
 
-    submit.click(on_submit, inputs=[input_video, model_type, remove_bg, max_c, max_d, lt, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
+    submit.click(on_submit, inputs=[input_video, model_type, remove_bg, max_c, max_d, max_s, max_l, max_v, lt, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
     render.click(None, inputs=[coords, mesh_order, bgcolor, output_frame, output_mask, selected, output_depth], outputs=None, js=load_model)
     render.click(partial(get_mesh), inputs=[output_frame, output_mask, blur_in, load_all], outputs=[result, result_file, mesh_order])
 
-    example_files = [["./examples/streetview.mp4", "vits", False, 12, 12, "slider", example_coords], ["./examples/man-in-museum-reverse-cut.mp4", "vits", True, 12, 12, "slider", example_coords]]
-    examples = gr.Examples(examples=example_files, fn=on_submit, cache_examples=True, inputs=[input_video, model_type, remove_bg, max_c, max_d, lt, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
+    example_files = [["./examples/streetview.mp4", "vits", False, 12, 12, 32, 64, 16, "slider", example_coords], ["./examples/man-in-museum-reverse-cut.mp4", "vits", True, 12, 12, 32, 64, 16, "slider", example_coords]]
+    examples = gr.Examples(examples=example_files, fn=on_submit, cache_examples=True, inputs=[input_video, model_type, remove_bg, max_c, max_d, max_s, max_l, max_v, lt, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
     
 
 if __name__ == '__main__':