app.py

import gradio as gr
import cv2
import os
import random
import zipfile
from PIL import Image
from datetime import datetime

def procesar_video(video_path):
    try:
        # Configurar directorio temporal
        temp_dir = f"temp_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        os.makedirs(temp_dir, exist_ok=True)
        
        # Leer el video y extraer TODOS los frames
        cap = cv2.VideoCapture(video_path)
        frame_count = 0
        frame_paths = []
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # Guardar cada frame
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)
            img_path = os.path.join(temp_dir, f"frame_{frame_count:04d}.jpg")
            img.save(img_path)
            frame_paths.append(img_path)
            frame_count += 1
        
        cap.release()
        
        if frame_count == 0:
            raise ValueError("No se pudieron extraer frames del video")
        
        # Seleccionar 4 frames aleatorios para el collage
        selected_frames = random.sample(frame_paths, min(4, len(frame_paths)))
        
        # Crear collage
        images = [Image.open(img) for img in selected_frames]
        collage = Image.new('RGB', (images[0].width * 2, images[0].height * 2))
        
        for i, img in enumerate(images):
            row = i // 2
            col = i % 2
            x = col * images[0].width
            y = row * images[0].height
            collage.paste(img, (x, y))
        
        collage_path = os.path.join(temp_dir, "collage.jpg")
        collage.save(collage_path)
        
        # Crear archivo ZIP con TODOS los frames
        zip_path = os.path.join(temp_dir, "frames.zip")
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for img in frame_paths:
                zipf.write(img, os.path.basename(img))
        
        return collage_path, zip_path, temp_dir
    
    except Exception as e:
        raise gr.Error(f"Error al procesar el video: {str(e)}")

def limpiar_cache(temp_dir):
    if temp_dir and os.path.exists(temp_dir):
        for file in os.listdir(temp_dir):
            os.remove(os.path.join(temp_dir, file))
        os.rmdir(temp_dir)

with gr.Blocks(title="Extracción de Fotogramas Forenses") as demo:
    gr.Markdown("# Herramienta de Extracción de Fotogramas Forenses")
    gr.Markdown("**Carga un video para extraer TODOS los fotogramas y generar un collage de muestra.**")
    gr.Markdown("Desarrollado por José R. Leonett para el Grupo de Peritos Forenses Digitales de Guatemala - [www.forensedigital.gt](https://www.forensedigital.gt)")
    
    with gr.Row():
        with gr.Column():
            video_input = gr.Video(label="Subir Video", interactive=True)
            procesar_btn = gr.Button("Procesar Fotogramas", interactive=False)
        
        with gr.Column():
            gallery_output = gr.Image(label="Collage de Muestra")
            download_btn = gr.Button("DESCARGAR FOTOGRAMAS", interactive=False)
            download_file = gr.File(label="Archivo ZIP generado", visible=False)
    
    temp_dir_state = gr.State()
    
    def habilitar_procesar(video):
        return gr.Button(interactive=True) if video else gr.Button(interactive=False)
    
    def procesar_y_mostrar(video):
        limpiar_cache(temp_dir_state.value)
        collage_path, zip_path, temp_dir = procesar_video(video)
        return collage_path, zip_path, temp_dir, gr.Button(interactive=True)
    
    video_input.change(
        fn=habilitar_procesar,
        inputs=video_input,
        outputs=procesar_btn,
        queue=False
    )
    
    procesar_btn.click(
        fn=procesar_y_mostrar,
        inputs=video_input,
        outputs=[gallery_output, download_file, temp_dir_state, download_btn],
    )
    
    download_btn.click(
        fn=lambda temp_dir: os.path.join(temp_dir, "frames.zip"),
        inputs=temp_dir_state,
        outputs=download_file,
    )

if __name__ == "__main__":
    demo.launch()