Update app.py

app.py

@@ -73,7 +73,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', blur_data=blurin, o=1, b=32):
+def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data=blurin, o=1, b=32, ctr=True):
     if encoder not in ["vitl","vitb","vits","vitg"]:
         encoder = "vits"
 
@@ -241,7 +241,9 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data
             masks.append(f"f{count}_mask.png")            
             count += 1
 
-        color_transfer(0)
+        global frame_selected
+        if ctr === True:
+            color_transfer(frame_selected)
 
         j=0
         while j<len(frames):
@@ -681,6 +683,7 @@ with gr.Blocks(css=css, js=js, head=head) as demo:
               input_url = gr.Textbox(elem_id="url_in", value="./examples/streetview.mp4", label="URL")
               input_video = gr.Video(label="Input Video", format="mp4")
               input_url.input(fn=loadfile, inputs=[input_url], outputs=[input_video])
+              ct = gr.Checkbox(label="Color transfer", value=True)
               submit = gr.Button("Submit")
             with gr.Group():
               output_frame = gr.Gallery(label="Frames", preview=True, columns=8192, interactive=False)
@@ -782,7 +785,7 @@ with gr.Blocks(css=css, js=js, head=head) as demo:
           input_json.input(show_json, inputs=[input_json], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth])
             
     
-    def on_submit(uploaded_video,model_type,blur_in,boffset,bsize,coordinates):
+    def on_submit(uploaded_video,model_type,blur_in,boffset,bsize,ct,coordinates):
         global locations
         locations = []
         avg = [0, 0]
@@ -841,14 +844,14 @@ with gr.Blocks(css=css, js=js, head=head) as demo:
         f.close()
             
         # Process the video and get the path of the output video
-        output_video_path = make_video(uploaded_video,encoder=model_type,blur_data=blurin,o=boffset,b=bsize)
+        output_video_path = make_video(uploaded_video,encoder=model_type,blur_data=blurin,o=boffset,b=bsize,ctr=ct)
 
         return output_video_path
 
-    submit.click(on_submit, inputs=[input_video, model_type, blur_in, boffset, bsize, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth])
+    submit.click(on_submit, inputs=[input_video, model_type, blur_in, boffset, bsize, ct, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth])
 
-    example_files = [["./examples/streetview.mp4", "vits", blurin, 1, 32, example_coords]]
-    examples = gr.Examples(examples=example_files, fn=on_submit, cache_examples=True, inputs=[input_video, model_type, blur_in, boffset, bsize, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth])
+    example_files = [["./examples/streetview.mp4", "vits", blurin, 1, 32, True, example_coords]]
+    examples = gr.Examples(examples=example_files, fn=on_submit, cache_examples=True, inputs=[input_video, model_type, blur_in, boffset, bsize, ct, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth])
     
 
 if __name__ == '__main__':