Fotogramas / app.py
leonett's picture
Update app.py
af3b295 verified
raw
history blame
8.92 kB
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()