Update app.py

app.py

@@ -1,7 +1,8 @@
 # app.py — veureu/svision (Salamandra Vision 7B · ZeroGPU) — compatible con ENGINE
 import os
 import json
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Any
+import io
 
 import gradio as gr
 import spaces
@@ -11,6 +12,11 @@ import numpy as np
 from PIL import Image
 from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration
 
+import cv2
+from scenedetect import VideoManager, SceneManager
+from scenedetect.detectors import ContentDetector
+
+
 MODEL_ID = os.environ.get("MODEL_ID", "BSC-LT/salamandra-7b-vision")
 DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -131,6 +137,46 @@ def face_image_embedding(image: Image.Image) -> List[float] | None:
         print(f"Fallo embedding cara: {e}")
         return None
 
+@spaces.GPU
+def scenes_extraction(video_file: io.IOBase, threshold: float, offset_frames: int, crop_ratio: float) -> Tuple[List[Image.Image], List[Dict]] | None:
+    # Detectamos las escenas
+    video_path = video_file.name
+    video_manager = VideoManager([video_path])
+    scene_manager = SceneManager()
+    scene_manager.add_detector(ContentDetector(threshold=threshold))
+    video_manager.start()
+    scene_manager.detect_scenes(video_manager)
+    scene_list = scene_manager.get_scene_list()
+
+    cap = cv2.VideoCapture(video_path)    
+    images: List[Image.Image] = []
+    informacion_escenas: List[Dict] = []
+
+    for i, (start_time, end_time) in enumerate(scene_list):
+        frame_number = int(start_time.get_frames()) + offset_frames
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
+        ret, frame = cap.read()
+        if ret:
+            h, w = frame.shape[:2]
+
+            # Ahora realizamos el recorte
+            ch, cw = int(h * crop_ratio), int(w * crop_ratio)
+            frame = frame[ch:h-ch, cw:w-cw]
+            
+            # Guardamos la escena obtenida
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            img = Image.fromarray(frame_rgb)
+            images.append(img)
+            
+            # Guardamos la información de la escena
+            informacion_escenas.append({
+                "index": i+1,
+                "start": start_time.get_seconds(),
+                "end": end_time.get_seconds()
+            })
+
+    cap.release()
+    return images, informacion_escenas
 
 # ----------------------------- UI & Endpoints --------------------------------
 
@@ -179,6 +225,15 @@ with gr.Blocks(title="Salamandra Vision 7B · ZeroGPU") as demo:
         face_out = gr.JSON(label="Embedding facial (vector)")
     face_btn.click(face_image_embedding, [face_img], face_out, api_name="face_image_embedding", concurrency_limit=1)
 
+    with gr.Row():
+        video_file = gr.Video(label="Sube un vídeo", type="file")
+        threshold = gr.Slider(0.0, 100.0, value=30.0, step=1.0, label="Threshold")
+        offset_frames = gr.Slider(0, 30, value=5, step=1, label="Offset frames")
+        crop_ratio = gr.Slider(0.0, 1.0, value=1.0, step=0.05, label="Crop ratio")
+        scenes_btn = gr.Button("Extraer escenas")
+        scenes_gallery_out = gr.Gallery(label="Keyframes de escenas", show_label=False, columns=4, height="auto")
+        scenes_info_out = gr.JSON(label="Información de escenas")
+    scenes_btn.click(scenes_extraction, inputs=[video_file, threshold, offset_frames, crop_ratio], outputs=[scenes_gallery_out, scenes_info_out], api_name="scenes_extraction", concurrency_limit=1)
 
 demo.queue(max_size=16).launch()