Update app.py

app.py

@@ -80,7 +80,7 @@ def predict_depth(image, model):
 #    return model(image)["depth"]
 
 def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg=False):
-    if encoder not in ["vitl","vitb","vits"]:
+    if encoder not in ["vitl","vitb","vits","vitg"]:
         encoder = "vits"
 
     model_name = encoder2name[encoder]
@@ -210,9 +210,21 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
             #raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2BGRA)
             #raw_frame[:, :, 3] = 255
 
-            if remove_bg == True:
+            if remove_bg == False:
+                thumbnail = cv2.cvtColor(cv2.resize(raw_frame, (16,32)), cv2.COLOR_BGR2GRAY).flatten()
+                if len(thumbnail_old) > 0:
+                    diff = thumbnail - thumbnail_old
+                    #print(diff)
+                    c = Counter(diff)
+                    value, cc = c.most_common()[0]
+                    if value == 0 and cc > int(16*32*0.8):
+                        count += 1
+                        continue
+                thumbnail_old = thumbnail
+            else:
                 if count >= 1: #int(cframes/2):
                     n = 0 #n = count-int(cframes/2)
+                    
                     depth_color_bg = cv2.imread(f"f{n}_dmap.png").astype(np.uint8)
                     raw_frame_bg = cv2.imread(f"f{n}.png").astype(np.uint8)
                     
@@ -228,15 +240,15 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
                     print('-median')
                     print(md_d)
                     print(md_c)
-                    print('-most common')
-                    c = Counter(diff_d.flatten())
-                    value, cc = c.most_common()[0]
-                    print(value)
-                    print(cc)
-                    c = Counter(diff_c.flatten())
-                    value, cc = c.most_common()[0]
-                    print(value)
-                    print(cc)
+                    #print('-most common')
+                    #c = Counter(diff_d.flatten())
+                    #value, cc = c.most_common()[0]
+                    #print(value)
+                    #print(cc)
+                    #c = Counter(diff_c.flatten())
+                    #value, cc = c.most_common()[0]
+                    #print(value)
+                    #print(cc)
                     print('-')
                     
                     mask_bg = cv2.inRange(diff_d, np.array([0,0,0]), np.array([md_d,md_d,md_d]))
@@ -251,16 +263,6 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
                     raw_frame[cv2.bitwise_and(mask_shadow, mask_bg)>0] = (raw_frame[cv2.bitwise_and(mask_shadow, mask_bg)>0] / 17 + 240).astype(np.uint8)
                     raw_frame[cv2.bitwise_and(mask_no_shadow, mask_bg)>0] = (255,255,255)
             
-            thumbnail = cv2.cvtColor(cv2.resize(raw_frame, (16,32)), cv2.COLOR_BGR2GRAY).flatten()
-            if len(thumbnail_old) > 0:
-                diff = thumbnail - thumbnail_old
-                #print(diff)
-                c = Counter(diff)
-                value, cc = c.most_common()[0]
-                if value == 0 and cc > int(16*32*0.8):
-                    continue
-            thumbnail_old = thumbnail
-            
             cv2.imwrite(f"f{count}.png", raw_frame)
             orig_frames.append(f"f{count}.png")
             
@@ -1024,7 +1026,7 @@ with gr.Blocks(css=css, js=js) as demo:
             find_normals.click(fn=findNormals, inputs=[format_normals], outputs=[normals_out])
 
         with gr.Column():
-            model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')
+            model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl"), ("giant", "vitg")], type="value", value="vits", label='Model Type')
             remove_bg = gr.Checkbox(label="Remove background")
             processed_video = gr.Video(label="Output Video", format="mp4", interactive=False)
             processed_zip = gr.File(label="Output Archive", interactive=False)
@@ -1404,7 +1406,7 @@ with gr.Blocks(css=css, js=js) as demo:
     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, example_coords], ["./examples/man-in-museum-reverse.mp4", "vits", True, example_coords]]
+    example_files = [["./examples/streetview.mp4", "vits", False, example_coords], ["./examples/man-in-museum-reverse-cut.mp4", "vits", True, example_coords]]
     examples = gr.Examples(examples=example_files, fn=on_submit, cache_examples=True, inputs=[input_video, model_type, remove_bg, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])