Update app.py

app.py

@@ -10,7 +10,6 @@ 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"))
@@ -22,14 +21,11 @@ def procesar_video(video):
         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
+        # Extracción de fotogramas
         cap = cv2.VideoCapture(video_path)
         frame_count = 0
         frame_paths = []
@@ -48,13 +44,13 @@ def procesar_video(video):
         if frame_count == 0:
             raise gr.Error("No se pudieron extraer fotogramas del video")
 
-        # Crear collage de 4 fotogramas aleatorios con separación
+        # Crear collage
         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))
+        collage = Image.new('RGB', 
+            (images[0].width * 2 + margin, images[0].height * 2 + margin), 
+            (255,255,255))
         positions = [
             (0, 0),
             (images[0].width + margin, 0),
@@ -66,61 +62,22 @@ def procesar_video(video):
         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
+        # Generación del ZIP con todos los fotogramas
         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)
-        final_zip_path = os.path.abspath(final_zip_path)
-        # ===============================
         
+        with zipfile.ZipFile(final_zip_path, mode="w") as zipf:
+            for img_path in frame_paths:
+                zipf.write(img_path, os.path.basename(img_path))
+            
+            chain_content = (
+                f"Nombre del archivo: {original_name}\n"
+                f"Fecha de carga y extracción: {timestamp}\n"
+                f"Número de fotogramas: {frame_count}\n"
+            )
+            zipf.writestr("cadena_custodia.txt", chain_content)
+
         return collage_path, final_zip_path, temp_dir
 
     except Exception as e:
@@ -132,40 +89,27 @@ def limpiar_cache(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("# 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 (solo archivos de video)", interactive=True)
+            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)
     
-    # Estados para guardar el directorio temporal y la ruta del ZIP
     temp_dir_state = gr.State(None)
     zip_path_state = gr.State(None)
     
     def on_video_change(video):
-        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)")
-        return gr.update(interactive=True), gr.update(interactive=False)
+        if video:
+            return gr.update(interactive=True), gr.update(interactive=False)
+        return gr.update(interactive=False), gr.update(interactive=False)
     
     video_input.change(
         fn=on_video_change,
@@ -175,13 +119,10 @@ with gr.Blocks(title="Extracción de Fotogramas Forenses") as demo:
     )
     
     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)
-        # Se retorna zip_path tanto para actualizar el estado como para asignar el valor al componente de descarga
-        return collage_path, zip_path, temp_dir, zip_path, btn_download_update, btn_proc_update
+        return collage_path, zip_path, temp_dir, zip_path, gr.update(interactive=True), gr.update(interactive=False)
     
     procesar_btn.click(
         fn=procesar_y_mostrar,
@@ -190,10 +131,9 @@ with gr.Blocks(title="Extracción de Fotogramas Forenses") as demo:
     )
     
     def trigger_download(zip_path):
-        # Se abre el archivo ZIP, se leen sus bytes y se retorna un par (nombre, contenido)
         with open(zip_path, "rb") as f:
             file_bytes = f.read()
-        return (os.path.basename(zip_path), file_bytes)
+        return os.path.basename(zip_path), file_bytes
     
     download_btn.click(
         fn=trigger_download,
@@ -202,4 +142,4 @@ with gr.Blocks(title="Extracción de Fotogramas Forenses") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()