Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -65,70 +65,78 @@ def get_prompt(click_state, click_input):
 # extract frames from upload video
 
 def get_frames_from_video(video_input, video_state):
+    """
+    Args:
+        video_path:str
+        timestamp:float64
+    Return 
+        19 элементов для Gradio-интерфейса
+    """
     video_path = video_input
     frames = []
     user_name = time.time()
     status_ok = True
     operation_log = [("[Must Do]", "Click image"), (": Video uploaded! Try to click the image shown in step2 to add masks.\n", None)]
     try:
-        # Получаем FPS из видео
+        # Считываем видео с помощью ffmpeg
         probe = ffmpeg.probe(video_path)
-        video_streams = [stream for stream in probe['streams'] if stream['codec_type'] == 'video']
-        fps_str = video_streams[0]['r_frame_rate']  # Например: "25/1"
-        fps = eval(fps_str)
-
-        # Извлекаем кадры с максимальным качеством во временную папку
-        with tempfile.TemporaryDirectory() as tmpdir:
-            frame_pattern = os.path.join(tmpdir, 'frame_%05d.png')
-            (
-                ffmpeg
-                .input(video_path)
-                .output(frame_pattern, start_number=0, vsync=0, qscale=0)
-                .run(quiet=True)
-            )
-            extracted = sorted(os.listdir(tmpdir))
-            for file in extracted:
-                img = Image.open(os.path.join(tmpdir, file)).convert("RGB")
-                frames.append(np.array(img))
+        video_info_ff = next(stream for stream in probe['streams'] if stream['codec_type'] == 'video')
+        fps = eval(video_info_ff['r_frame_rate'])  # Например, '25/1'
+        width = int(video_info_ff['width'])
+        height = int(video_info_ff['height'])
+        total_frames = int(video_info_ff['nb_frames'])
 
-        original_h, original_w = frames[0].shape[:2]
-
-        if len(frames) >= 500:
+        if total_frames >= 500:
             operation_log = [("You uploaded a video with more than 500 frames. Stop the video extraction. Kindly lower the video frame rate to a value below 500. We highly recommend deploying the demo locally for long video processing.", "Error")]
             status_ok = False
+            return [None] * 19  # Возвращаем заглушки
+        else:
+            out, _ = (
+                ffmpeg
+                .input(video_path)
+                .output('pipe:', format='rawvideo', pix_fmt='rgb24')
+                .run(capture_stdout=True, capture_stderr=True)
+            )
+            video = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])
+            frames = [Image.fromarray(f) for f in video]
 
+            print(f'Inp video shape: t_{len(frames)}, s_{width}x{height}')
     except Exception as e:
         status_ok = False
         print("read_frame_source:{} error. {}\n".format(video_path, str(e)))
+        return [None] * 19
 
-    if frames and (frames[0].shape[0] > 720 or frames[0].shape[1] > 720):
-        operation_log = [(f"Video uploaded! Try to click the image shown in step2 to add masks. (You uploaded a video with a size of {original_w}x{original_h}, and the length of its longest edge exceeds 720 pixels. We may resize the input video during processing.)", "Normal")]
+    if frames[0].size[1] > 720 or frames[0].size[0] > 720:
+        operation_log = [(f"Video uploaded! Try to click the image shown in step2 to add masks. (You uploaded a video with a size of {width}x{height}, and the length of its longest edge exceeds 720 pixels. We may resize the input video during processing.)", "Normal")]
 
     video_state = {
         "user_name": user_name,
         "video_name": os.path.split(video_path)[-1],
-        "origin_images": frames,
-        "painted_images": frames.copy(),
-        "masks": [np.zeros((original_h, original_w), np.uint8)]*len(frames),
-        "logits": [None]*len(frames),
+        "origin_images": [np.array(f) for f in frames],
+        "painted_images": [np.array(f) for f in frames],
+        "masks": [np.zeros((height, width), np.uint8)] * len(frames),
+        "logits": [None] * len(frames),
         "select_frame_number": 0,
         "fps": fps
     }
 
     video_info = "Video Name: {},\nFPS: {},\nTotal Frames: {},\nImage Size:{}".format(
-        video_state["video_name"], round(video_state["fps"], 0), len(frames), (original_w, original_h)
+        video_state["video_name"], round(video_state["fps"], 0), len(frames), (width, height)
     )
-    model.samcontroler.sam_controler.reset_image()
+    model.samcontroler.sam_controler.reset_image() 
     model.samcontroler.sam_controler.set_image(video_state["origin_images"][0])
 
-    return video_state, video_info, video_state["origin_images"][0], gr.update(visible=status_ok, maximum=len(frames), value=1), gr.update(visible=status_ok, maximum=len(frames), value=len(frames)), \
-           gr.update(visible=status_ok), gr.update(visible=status_ok), \
-           gr.update(visible=status_ok), gr.update(visible=status_ok), \
-           gr.update(visible=status_ok), gr.update(visible=status_ok), \
-           gr.update(visible=status_ok), gr.update(visible=status_ok), \
-           gr.update(visible=status_ok), gr.update(visible=status_ok), \
-           gr.update(visible=status_ok, choices=[], value=[]), \
-           gr.update(visible=True, value=operation_log), gr.update(visible=status_ok, value=operation_log)
+    return video_state, video_info, frames[0], \
+        gr.update(visible=status_ok, maximum=len(frames), value=1), \
+        gr.update(visible=status_ok, maximum=len(frames), value=len(frames)), \
+        gr.update(visible=status_ok), gr.update(visible=status_ok), \
+        gr.update(visible=status_ok), gr.update(visible=status_ok), \
+        gr.update(visible=status_ok), gr.update(visible=status_ok), \
+        gr.update(visible=status_ok), gr.update(visible=status_ok), \
+        gr.update(visible=status_ok), gr.update(visible=status_ok), \
+        gr.update(visible=status_ok), gr.update(visible=status_ok, choices=[], value=[]), \
+        gr.update(visible=True, value=operation_log), \
+        gr.update(visible=status_ok, value=operation_log)
 
 # get the select frame from gradio slider
 def select_template(image_selection_slider, video_state, interactive_state, mask_dropdown):