Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -16,6 +16,7 @@ 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
@@ -66,91 +67,95 @@ 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):
     """
-    Извлечение кадров без потерь через ffmpeg (16-битные RGB кадры).
+    Извлечение кадров из видео без потерь (16-битный RGB).
     """
-    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)]
 
-    output_dir = f"/tmp/frames_{user_name}"
-    os.makedirs(output_dir, exist_ok=True)
-    output_pattern = os.path.join(output_dir, "frame_%05d.tiff")
-
     try:
-        # Извлекаем 16-битные RGB кадры
-        cmd = [
-            "ffmpeg", "-y", "-i", video_path,
-            "-pix_fmt", "rgb48le",  # 16-бит на канал
-            output_pattern
+        # Папка под кадры
+        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(cmd, check=True)
+        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.")
 
-        # Загружаем изображения
-        frame_files = sorted(f for f in os.listdir(output_dir) if f.endswith(".tiff"))
-        for frame_name in frame_files:
-            image = Image.open(os.path.join(output_dir, frame_name))
+        for path in frame_paths:
+            image = Image.open(path)
             frames.append(np.array(image))
 
-        if not frames:
-            raise RuntimeError("Не удалось извлечь кадры.")
+        # Получение 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)
 
-        # Определим fps через ffprobe
-        try:
-            result = subprocess.run(
-                ["ffprobe", "-v", "error", "-select_streams", "v:0",
-                 "-show_entries", "stream=r_frame_rate", "-of", "default=noprint_wrappers=1:nokey=1",
-                 video_path],
-                stdout=subprocess.PIPE, stderr=subprocess.STDOUT
-            )
-            fps_expr = result.stdout.decode().strip()
-            num, denom = map(int, fps_expr.split('/'))
-            fps = num / denom
-        except:
-            fps = 24  # fallback
-
-        if len(frames) >= 600:
-            operation_log = [("Видео содержит более 500 кадров. Обработка остановлена. Запустите локально.", "Error")]
+        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}")
         status_ok = False
-        operation_log = [(f"[Ошибка] Не удалось извлечь кадры: {e}", "Error")]
-        return None, None, None, *[gr.update(visible=False)] * 16, gr.update(visible=True, value=operation_log), gr.update(visible=False, value=operation_log)
-
-    if frames[0].shape[0] > 720 or frames[0].shape[1] > 720:
-        operation_log = [(f"Видео размером {original_w}x{original_h}. Возможно будет ресайз при обработке.", "Normal")]
+        return None
 
+    # Собираем состояние
     video_state = {
         "user_name": user_name,
-        "video_name": os.path.basename(video_path),
+        "video_name": os.path.basename(video_input),
         "origin_images": frames,
         "painted_images": frames.copy(),
-        "masks": [np.zeros((original_h, original_w), np.uint8)] * len(frames),
-        "logits": [None] * len(frames),
+        "masks": [np.zeros((original_h, original_w), np.uint8)] * length,
+        "logits": [None] * length,
         "select_frame_number": 0,
         "fps": fps
     }
 
-    video_info = f"Video Name: {video_state['video_name']},\nFPS: {round(fps, 0)},\nTotal Frames: {len(frames)},\nImage Size: {original_w}x{original_h}"
+    video_info = f"Video Name: {video_state['video_name']},\nFPS: {round(video_state['fps'], 0)},\nTotal Frames: {length},\nImage Size: {(original_w, original_h)}"
 
-    # SAM-подгрузка — не удаляй, если используется
     model.samcontroler.sam_controler.reset_image()
-    model.samcontroler.sam_controler.set_image(frames[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) for _ in range(14)], \
-           gr.update(visible=status_ok, choices=[], value=[]), \
-           gr.update(visible=True, value=operation_log), \
-           gr.update(visible=status_ok, value=operation_log)
-
+    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)
 
 def select_template(image_selection_slider, video_state, interactive_state, mask_dropdown):