Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -70,7 +70,7 @@ def get_frames_from_video(video_input, video_state):
         video_path:str
         timestamp:float64
     Return 
-        19 элементов для Gradio-интерфейса
+        [[0:nearest_frame], [nearest_frame:], nearest_frame]
     """
     video_path = video_input
     frames = []
@@ -78,65 +78,60 @@ def get_frames_from_video(video_input, video_state):
     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:
-        # Считываем видео с помощью ffmpeg
-        probe = ffmpeg.probe(video_path)
-        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'])
-
-        if total_frames >= 500:
+        cap = cv2.VideoCapture(video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+        if length >= 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
-            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:
+            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:
         status_ok = False
         print("read_frame_source:{} error. {}\n".format(video_path, str(e)))
-        return [None] * 19
-
-    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")]
+    
+    # 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")]
 
     video_state = {
         "user_name": user_name,
         "video_name": os.path.split(video_path)[-1],
-        "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),
+        "origin_images": frames,
+        "painted_images": frames.copy(),
+        "masks": [np.zeros((original_h, original_w), 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), (width, height)
-    )
+        }
+    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() 
     model.samcontroler.sam_controler.set_image(video_state["origin_images"][0])
-
-    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)
+    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)
 
 # get the select frame from gradio slider
 def select_template(image_selection_slider, video_state, interactive_state, mask_dropdown):