Spaces:
Running
Running
import gradio as gr | |
import cv2 | |
import os | |
import random | |
import zipfile | |
from PIL import Image | |
from datetime import datetime | |
import io | |
import hashlib | |
def procesar_video(video): | |
try: | |
# Validar y obtener el nombre y ruta del video | |
if isinstance(video, dict): | |
original_name = video.get("name", "video") | |
video_path = video.get("file", video.get("data")) | |
else: | |
original_name = os.path.basename(video) | |
video_path = video | |
allowed_extensions = ('.mp4', '.avi', '.mov', '.mkv') | |
if not original_name.lower().endswith(allowed_extensions): | |
raise gr.Error("Solo se permiten archivos de video (mp4, avi, mov, mkv)") | |
# Timestamp para la cadena de custodia | |
timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S') | |
# Crear directorio temporal | |
temp_dir = f"temp_{datetime.now().strftime('%Y%m%d%H%M%S')}" | |
os.makedirs(temp_dir, exist_ok=True) | |
# Extraer TODOS los fotogramas del video | |
cap = cv2.VideoCapture(video_path) | |
frame_count = 0 | |
frame_paths = [] | |
while True: | |
ret, frame = cap.read() | |
if not ret: | |
break | |
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 gr.Error("No se pudieron extraer fotogramas del video") | |
# Crear collage de 4 fotogramas aleatorios con separaci贸n | |
selected_frames = random.sample(frame_paths, min(4, len(frame_paths))) | |
images = [Image.open(img) for img in selected_frames] | |
margin = 10 | |
collage_width = images[0].width * 2 + margin | |
collage_height = images[0].height * 2 + margin | |
collage = Image.new('RGB', (collage_width, collage_height), color=(255,255,255)) | |
positions = [ | |
(0, 0), | |
(images[0].width + margin, 0), | |
(0, images[0].height + margin), | |
(images[0].width + margin, images[0].height + margin) | |
] | |
for i, img in enumerate(images): | |
collage.paste(img, positions[i]) | |
collage_path = os.path.join(temp_dir, "collage.jpg") | |
collage.save(collage_path) | |
# =============================== | |
# Generaci贸n de la cadena de custodia | |
# ------------------------------- | |
# 1. Crear ZIP provisional en memoria con SOLO los fotogramas | |
provisional_zip = io.BytesIO() | |
with zipfile.ZipFile(provisional_zip, mode="w") as zipf: | |
for img in frame_paths: | |
zipf.write(img, os.path.basename(img)) | |
provisional_zip.seek(0) | |
provisional_data = provisional_zip.getvalue() | |
md5_provisional = hashlib.md5(provisional_data).hexdigest() | |
# 2. Crear un ZIP provisional en memoria que incluya un archivo TXT con marcador de posici贸n para el MD5 final | |
final_zip_mem = io.BytesIO() | |
chain_placeholder = "MD5 final del ZIP: <pending>" | |
chain_initial = ( | |
f"Nombre del archivo: {original_name}\n" | |
f"Fecha de carga y extracci贸n: {timestamp}\n" | |
f"N煤mero de fotogramas: {frame_count}\n" | |
f"MD5 provisional (frames): {md5_provisional}\n" | |
f"{chain_placeholder}\n" | |
) | |
with zipfile.ZipFile(final_zip_mem, mode="w") as zipf: | |
for img in frame_paths: | |
zipf.write(img, os.path.basename(img)) | |
zipf.writestr("chain_of_custody.txt", chain_initial) | |
final_zip_mem.seek(0) | |
final_zip_data = final_zip_mem.getvalue() | |
md5_final = hashlib.md5(final_zip_data).hexdigest() | |
# 3. Reconstruir el ZIP final en memoria con el TXT actualizado (incluyendo ambos MD5) | |
final_chain = ( | |
f"Nombre del archivo: {original_name}\n" | |
f"Fecha de carga y extracci贸n: {timestamp}\n" | |
f"N煤mero de fotogramas: {frame_count}\n" | |
f"MD5 provisional (frames): {md5_provisional}\n" | |
f"MD5 final del ZIP: {md5_final}\n" | |
) | |
final_zip_mem2 = io.BytesIO() | |
with zipfile.ZipFile(final_zip_mem2, mode="w") as zipf: | |
for img in frame_paths: | |
zipf.write(img, os.path.basename(img)) | |
zipf.writestr("chain_of_custody.txt", final_chain) | |
final_zip_mem2.seek(0) | |
final_zip_data2 = final_zip_mem2.getvalue() | |
# Escribir el ZIP final a disco con el nombre basado en el video original | |
base_name = os.path.splitext(original_name)[0] | |
zip_filename = f"{base_name}.zip" | |
final_zip_path = os.path.join(temp_dir, zip_filename) | |
with open(final_zip_path, "wb") as f: | |
f.write(final_zip_data2) | |
# =============================== | |
return collage_path, final_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) | |
# ------------------------------- | |
# Interfaz Gradio | |
# ------------------------------- | |
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(): | |
# Desactivar opciones de webcam y edici贸n (tijera) estableciendo source="upload" | |
video_input = gr.Video(source="upload", label="Subir Video (solo archivos de video)", interactive=True) | |
# Ambos botones inician desactivados | |
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) | |
# Estados para guardar el directorio temporal y la ruta del ZIP | |
temp_dir_state = gr.State(None) | |
zip_path_state = gr.State(None) | |
# Al cargar un video: si es v谩lido, se activa el bot贸n "Procesar Fotogramas" y se desactiva "DESCARGAR FOTOGRAMAS" | |
def on_video_change(video): | |
# Verificar que se haya cargado algo y que contenga informaci贸n | |
if not video: | |
return gr.update(interactive=False), gr.update(interactive=False) | |
if isinstance(video, dict) and video.get("name"): | |
original_name = video["name"] | |
elif isinstance(video, str) and video: | |
original_name = os.path.basename(video) | |
else: | |
return gr.update(interactive=False), gr.update(interactive=False) | |
allowed_extensions = ('.mp4', '.avi', '.mov', '.mkv') | |
if not original_name.lower().endswith(allowed_extensions): | |
raise gr.Error("Solo se permiten archivos de video (mp4, avi, mov, mkv)") | |
# Al tener un video v谩lido, activa "Procesar Fotogramas" y desactiva "DESCARGAR FOTOGRAMAS" | |
return gr.update(interactive=True), gr.update(interactive=False) | |
video_input.change( | |
fn=on_video_change, | |
inputs=video_input, | |
outputs=[procesar_btn, download_btn], | |
queue=False | |
) | |
# Al pulsar "Procesar Fotogramas": se procesa el video, se genera el collage y el ZIP, se desactiva el bot贸n de procesar | |
# y se activa el bot贸n de descarga. | |
def procesar_y_mostrar(video): | |
btn_proc_update = gr.update(interactive=False) | |
if temp_dir_state.value: | |
limpiar_cache(temp_dir_state.value) | |
collage_path, zip_path, temp_dir = procesar_video(video) | |
btn_download_update = gr.update(interactive=True) | |
return collage_path, zip_path, temp_dir, zip_path, btn_download_update, btn_proc_update | |
procesar_btn.click( | |
fn=procesar_y_mostrar, | |
inputs=video_input, | |
outputs=[gallery_output, download_file, temp_dir_state, zip_path_state, download_btn, procesar_btn], | |
) | |
# Funci贸n para activar la descarga: actualiza el componente File con la ruta del ZIP y lo hace visible | |
def trigger_download(zip_path): | |
return gr.update(value=zip_path, visible=True) | |
download_btn.click( | |
fn=trigger_download, | |
inputs=zip_path_state, | |
outputs=download_file, | |
) | |
if __name__ == "__main__": | |
demo.launch() | |