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
        cap = cv2.VideoCapture(video_path)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        # Seleccionar 4 fotogramas aleatorios
        frame_indices = sorted(random.sample(range(total_frames), 4))
        selected_frames = []
        
        for idx in frame_indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
            ret, frame = cap.read()
            if ret:
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img = Image.fromarray(frame_rgb)
                img_path = os.path.join(temp_dir, f"frame_{idx}.jpg")
                img.save(img_path)
                selected_frames.append(img_path)
        
        cap.release()
        
        # Crear collage
        images = [Image.open(img) for img in selected_frames]
        collage = Image.new('RGB', (images[0].width * 2, images[0].height * 2))
        
        collage.paste(images[0], (0, 0))
        collage.paste(images[1], (images[0].width, 0))
        collage.paste(images[2], (0, images[0].height))
        collage.paste(images[3], (images[0].width, images[0].height))
        
        collage_path = os.path.join(temp_dir, "collage.jpg")
        collage.save(collage_path)
        
        # Crear archivo ZIP
        zip_path = os.path.join(temp_dir, "frames.zip")
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for img in selected_frames:
                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 4 fotogramas aleatorios y generar un collage.**")
    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 Fotogramas")
            download_btn = gr.Button("DESCARGAR FOTOGRAMAS", interactive=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)
    
    def actualizar_descarga(temp_dir):
        return gr.Button(interactive=bool(temp_dir))
    
    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_btn, temp_dir_state, download_btn],
    )
    
    download_btn.click(
        fn=lambda temp_dir: os.path.join(temp_dir, "frames.zip"),
        inputs=temp_dir_state,
        outputs=gr.File(label="Frames descargados"),
    )

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