Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -2,7 +2,6 @@ import sys
 sys.path.append("../../")
 
 import os
-import subprocess
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 import json
 import time
@@ -16,7 +15,6 @@ import numpy as np
 import gradio as gr
 from gradio import Brush
 import tempfile
-import glob
 import ffmpeg
 from PIL import Image
 from tools.painter import mask_painter
@@ -67,95 +65,75 @@ def get_prompt(click_state, click_input):
 
 # extract frames from upload video
 
-import os
-import time
-import subprocess
-import glob
-import numpy as np
-from PIL import Image
-import gradio as gr
-
 def get_frames_from_video(video_input, video_state):
     """
-    Извлечение кадров из видео без потерь (16-битный RGB).
+    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:
-        # Папка под кадры
-        temp_output_dir = f"/tmp/frames_{user_name}"
-        os.makedirs(temp_output_dir, exist_ok=True)
-
-        # Извлечение кадров без сжатия (TIFF 16 бит, RGB)
-        command = [
-            "ffmpeg", "-i", video_input,
-            "-pix_fmt", "rgb48le",
-            os.path.join(temp_output_dir, "frame_%05d.tiff")
-        ]
-        subprocess.run(command, check=True)
-
-        # Получаем пути к кадрам
-        frame_paths = sorted(glob.glob(os.path.join(temp_output_dir, "frame_*.tiff")))
-        if len(frame_paths) == 0:
-            raise RuntimeError("ffmpeg failed to extract frames or no frames found.")
-
-        for path in frame_paths:
-            image = Image.open(path)
-            frames.append(np.array(image))
-
-        # Получение FPS через ffprobe
-        fps_cmd = [
-            "ffprobe", "-v", "error", "-select_streams", "v:0",
-            "-show_entries", "stream=r_frame_rate",
-            "-of", "default=noprint_wrappers=1:nokey=1", video_input
-        ]
-        fps_raw = subprocess.check_output(fps_cmd).decode().strip()
-        num, den = map(int, fps_raw.split("/")) if "/" in fps_raw else (int(fps_raw), 1)
-        fps = num / den
-
-        original_h, original_w = frames[0].shape[:2]
-        length = len(frames)
+        cap = cv2.VideoCapture(video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
 
         if length >= 600:
-            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.", "Error")]
-            status_ok = False
-
-    except Exception as e:
-        print(f"[FFmpeg ERROR]: {e}")
+            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:
         status_ok = False
-        return None
+        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")]
 
-    # Собираем состояние
     video_state = {
         "user_name": user_name,
-        "video_name": os.path.basename(video_input),
+        "video_name": os.path.split(video_path)[-1],
         "origin_images": frames,
         "painted_images": frames.copy(),
-        "masks": [np.zeros((original_h, original_w), np.uint8)] * length,
-        "logits": [None] * length,
+        "masks": [np.zeros((original_h, original_w), np.uint8)]*len(frames),
+        "logits": [None]*len(frames),
         "select_frame_number": 0,
         "fps": fps
-    }
-
-    video_info = f"Video Name: {video_state['video_name']},\nFPS: {round(video_state['fps'], 0)},\nTotal Frames: {length},\nImage Size: {(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), 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, video_state["origin_images"][0], \
-        gr.update(visible=status_ok, maximum=length, value=1), \
-        gr.update(visible=status_ok, maximum=length, value=length), \
-        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)
 
 def select_template(image_selection_slider, video_state, interactive_state, mask_dropdown):