Update app.py

app.py

@@ -1,37 +1,45 @@
+import av
 import torch
 import gradio as gr
 
 
+def get_video_length_av(video_path):
+    with av.open(video_path) as container:
+        stream = container.streams.video[0]
+        if container.duration is not None:
+            duration_in_seconds = float(container.duration) / av.time_base
+        else:
+            duration_in_seconds = stream.duration * stream.time_base
+
+    return duration_in_seconds
+
+
+def get_video_dimensions(video_path):
+    with av.open(video_path) as container:
+        video_stream = container.streams.video[0]
+        width = video_stream.width
+        height = video_stream.height
+
+    return width, height
+
+
 def get_free_memory_gb():
     gpu_index = torch.cuda.current_device()
-    # Get the GPU's properties
     gpu_properties = torch.cuda.get_device_properties(gpu_index)
 
-    # Get the total and allocated memory
     total_memory = gpu_properties.total_memory
     allocated_memory = torch.cuda.memory_allocated(gpu_index)
 
-    # Calculate the free memory
     free_memory = total_memory - allocated_memory
     return free_memory / 1024**3
 
 
-model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3")
-
-if torch.cuda.is_available():
-    free_memory = get_free_memory_gb()
-    concurrency_count = int(free_memory // 7)
-    model = model.cuda()
-    print(f"Using GPU with concurrency: {concurrency_count}")
-    print(f"Available video memory: {free_memory} GB")
-else:
-    print("Using CPU")
-    concurrency_count = 1
-
-convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")
-
-
 def inference(video):
+    if get_video_length_av(video.name) > 30:
+        raise gr.Error("Length of video cannot be over 30 seconds")
+    if get_video_dimensions(video.name) > (1920, 1920):
+        raise gr.Error("Video resolution must not be higher than 1920x1080")
+
     convert_video(
         model,  # The loaded model, can be on any device (cpu or cuda).
         input_source=video,  # A video file or an image sequence directory.
@@ -48,23 +56,40 @@ def inference(video):
     return "com.mp4"
 
 
-with gr.Blocks(title="Robust Video Matting") as block:
-    gr.Markdown("# Robust Video Matting")
-    gr.Markdown(
-        "Gradio demo for Robust Video Matting. To use it, simply upload your video, or click one of the examples to load them. Read more at the links below."
-    )
-    with gr.Row():
-        inp = gr.Video(label="Input Video")
-        out = gr.Video(label="Output Video")
-    btn = gr.Button("Run")
-    btn.click(inference, inputs=inp, outputs=out)
-
-    gr.Examples(
-        examples=[["example.mp4"]],
-        inputs=[inp],
-    )
-    gr.HTML(
-        "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.11515'>Robust High-Resolution Video Matting with Temporal Guidance</a> | <a href='https://github.com/PeterL1n/RobustVideoMatting'>Github Repo</a></p>"
-    )
+if __name__ == "__main__":
+    model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3")
+
+    if torch.cuda.is_available():
+        free_memory = get_free_memory_gb()
+        concurrency_count = int(free_memory // 7)
+        model = model.cuda()
+        print(f"Using GPU with concurrency: {concurrency_count}")
+        print(f"Available video memory: {free_memory} GB")
+    else:
+        print("Using CPU")
+        concurrency_count = 1
+
+    convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")
+
+    with gr.Blocks(title="Robust Video Matting") as block:
+        gr.Markdown("# Robust Video Matting")
+        gr.Markdown(
+            "Gradio demo for Robust Video Matting. To use it, simply upload your video, or click one of the examples to load them. Read more at the links below."
+        )
+        with gr.Row():
+            inp = gr.Video(label="Input Video")
+            out = gr.Video(label="Output Video")
+        btn = gr.Button("Run")
+        btn.click(inference, inputs=inp, outputs=out)
+
+        gr.Examples(
+            examples=[["example.mp4"]],
+            inputs=[inp],
+        )
+        gr.HTML(
+            "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.11515'>Robust High-Resolution Video Matting with Temporal Guidance</a> | <a href='https://github.com/PeterL1n/RobustVideoMatting'>Github Repo</a></p>"
+        )
 
-block.queue(api_open=False, max_size=5, concurrency_count=concurrency_count).launch()
+    block.queue(
+        api_open=False, max_size=5, concurrency_count=concurrency_count
+    ).launch()