Update app.py

app.py

@@ -9,12 +9,19 @@ import torch.nn.functional as F
 from torchvision import transforms
 from torchvision.transforms import Compose
 import tempfile
-import spaces 
+import spaces
+from zipfile import ZipFile
 
 from depth_anything.dpt import DepthAnything
 from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
 from moviepy.editor import *
 
+def zip_files(files):
+    with ZipFile("depth_result.zip", "w") as zipObj:
+        for idx, file in enumerate(files):
+            zipObj.write(file.name, file.name.split("/")[-1])
+    return "depth_result.zip"
+
 def create_video(frames, fps, type):
     print("building video result")
     clip = ImageSequenceClip(frames, fps=fps)
@@ -141,11 +148,12 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits'):
             count += 1
 
         final_vid = create_video(depth_frames, frame_rate, "depth")
+        final_zip = zip_files(depth_frames)
         raw_video.release()
         # out.release()
         cv2.destroyAllWindows()
         
-        return final_vid #output_path
+        return final_vid, final_zip #output_path
 
 def loadurl(url):
     return url
@@ -200,7 +208,8 @@ with gr.Blocks(css=css) as demo:
             submit = gr.Button("Submit")
         with gr.Column():
             model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')
-            processed_video = gr.Video(label="Processed Video", format="mp4")
+            processed_video = gr.Video(label="Output Video", format="mp4")
+            processed_zip = gr.File(label="Output Archive")
     
     def on_submit(uploaded_video,model_type):
                 
@@ -209,12 +218,12 @@ with gr.Blocks(css=css) as demo:
 
         return output_video_path
 
-    submit.click(on_submit, inputs=[input_video, model_type], outputs=processed_video)
+    submit.click(on_submit, inputs=[input_video, model_type], outputs=[processed_video, processed_zip])
 
     example_files = os.listdir('examples')
     example_files.sort()
     example_files = [os.path.join('examples', filename) for filename in example_files]
-    examples = gr.Examples(examples=example_files, inputs=[input_video], outputs=processed_video, fn=on_submit, cache_examples=True)
+    examples = gr.Examples(examples=example_files, inputs=[input_video], outputs=[processed_video, processed_zip], fn=on_submit, cache_examples=True)
     
 
 if __name__ == '__main__':