Update web-demos/hugging_face/app.py

web-demos/hugging_face/app.py

@@ -2,6 +2,7 @@ import sys
 sys.path.append("../../")
 
 import os
+import subprocess
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 import json
 import time
@@ -67,72 +68,89 @@ def get_prompt(click_state, click_input):
 
 def get_frames_from_video(video_input, video_state):
     """
-    Args:
-        video_path:str
-        timestamp:float64
-    Return 
-        [[0:nearest_frame], [nearest_frame:], nearest_frame]
+    Извлечение кадров без потерь через ffmpeg (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:
-        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. 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))
+        # Извлекаем 16-битные RGB кадры
+        cmd = [
+            "ffmpeg", "-y", "-i", video_path,
+            "-pix_fmt", "rgb48le",  # 16-бит на канал
+            output_pattern
+        ]
+        subprocess.run(cmd, check=True)
+
+        # Загружаем изображения
+        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))
+            frames.append(np.array(image))
+
+        if not frames:
+            raise RuntimeError("Не удалось извлечь кадры.")
+
+        original_h, original_w = frames[0].shape[:2]
+
+        # Определим 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")]
             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
+        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"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")]
+        operation_log = [(f"Видео размером {original_w}x{original_h}. Возможно будет ресайз при обработке.", "Normal")]
 
     video_state = {
         "user_name": user_name,
-        "video_name": os.path.split(video_path)[-1],
+        "video_name": os.path.basename(video_path),
         "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)] * 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), 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=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)
+    }
+
+    video_info = f"Video Name: {video_state['video_name']},\nFPS: {round(fps, 0)},\nTotal Frames: {len(frames)},\nImage Size: {original_w}x{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)
+
 
 def select_template(image_selection_slider, video_state, interactive_state, mask_dropdown):