Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -61,52 +61,48 @@ def get_prompt(click_state, click_input):
     return prompt
 
 # extract frames from upload video
+import tempfile
+import ffmpeg
+from PIL import Image
 def get_frames_from_video(video_input, video_state):
-    """
-    Args:
-        video_path:str
-        timestamp:float64
-    Return 
-        [[0:nearest_frame], [nearest_frame:], nearest_frame]
-    """
     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:
-        cap = cv2.VideoCapture(video_path)
-        fps = cap.get(cv2.CAP_PROP_FPS)
-        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-
-        if length >= 500:
+        # Получаем FPS из видео
+        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))
+
+        original_h, original_w = frames[0].shape[:2]
+
+        if len(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")]
-            ret, frame = cap.read()
-            if ret == True:
-                original_h, original_w = frame.shape[:2]
-                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
             status_ok = False
-        else:
-            while cap.isOpened():
-                ret, frame = cap.read()
-                if ret == True:
-                    # resize input image
-                    original_h, original_w = frame.shape[:2]
-                    frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-                else:
-                    break
-            t = len(frames)
-            if t > 0:
-                print(f'Inp video shape: t_{t}, s_{original_h}x{original_w}')
-            else:
-                print(f'Inp video shape: t_{t}, no input video!!!')
-    except (OSError, TypeError, ValueError, KeyError, SyntaxError) as e:
+
+    except Exception as e:
         status_ok = False
         print("read_frame_source:{} error. {}\n".format(video_path, str(e)))
-    
-    # initialize video_state
-    if 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 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")]
 
     video_state = {
         "user_name": user_name,
@@ -117,18 +113,22 @@ def get_frames_from_video(video_input, video_state):
         "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), length, (original_w, original_h))
-    model.samcontroler.sam_controler.reset_image() 
+    }
+
+    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)
+    )
+    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), gr.update(visible=status_ok, choices=[], value=[]), \
-                        gr.update(visible=True, value=operation_log), gr.update(visible=status_ok, value=operation_log)
+           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):
@@ -357,8 +357,7 @@ def inpaint_video(video_state, *_args):
     video_output = generate_video_from_frames(
         inpainted_all,
         output_path=output_path,
-        fps=fps,
-        bitrate="30M"
+        fps=fps
     )
 
     return video_output, operation_log, operation_log