File size: 3,987 Bytes
ee2a1ac
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
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()