Update app.py

app.py

@@ -6,6 +6,7 @@ import io
 import random
 import os
 import uuid
+import json # <-- Import the json library
 
 # --- Configuration & Setup ---
 # Create a directory to store temporary output files
@@ -17,18 +18,15 @@ os.makedirs(TEMP_DIR, exist_ok=True)
 # ... (All the functions like scramble_image, unscramble_image, etc., are the same)
 
 def process_image_for_grid(image_pil, grid_size):
-    """Crops the image to be perfectly divisible by the grid size."""
     if image_pil is None: return None, 0, 0
     img_width, img_height = image_pil.size
-    tile_w = img_width // grid_size
-    tile_h = img_height // grid_size
+    tile_w, tile_h = img_width // grid_size, img_height // grid_size
     if tile_w == 0 or tile_h == 0:
-        raise gr.Error(f"Image is too small for a {grid_size}x{grid_size} grid. Please use a larger image or a smaller grid size.")
+        raise gr.Error(f"Image is too small for a {grid_size}x{grid_size} grid.")
     cropped_image = image_pil.crop((0, 0, tile_w * grid_size, tile_h * grid_size))
     return cropped_image, tile_w, tile_h
 
 def scramble_image(image_pil, grid_size, seed):
-    """Scrambles an image and returns the PIL image and the map."""
     cropped_image, tile_w, tile_h = process_image_for_grid(image_pil, grid_size)
     tiles = [cropped_image.crop((j * tile_w, i * tile_h, (j + 1) * tile_w, (i + 1) * tile_h)) for i in range(grid_size) for j in range(grid_size)]
     rng = np.random.default_rng(seed=int(seed))
@@ -40,7 +38,6 @@ def scramble_image(image_pil, grid_size, seed):
     return scrambled_image, scramble_map
 
 def unscramble_image(scrambled_pil, scramble_map, grid_size):
-    """Unscrambles an image using the provided scramble_map."""
     cropped_image, tile_w, tile_h = process_image_for_grid(scrambled_pil, grid_size)
     scrambled_tiles = [cropped_image.crop((j * tile_w, i * tile_h, (j + 1) * tile_w, (i + 1) * tile_h)) for i in range(grid_size) for j in range(grid_size)]
     unscrambled_image = Image.new('RGB', cropped_image.size)
@@ -50,7 +47,6 @@ def unscramble_image(scrambled_pil, scramble_map, grid_size):
     return unscrambled_image
 
 def create_mapping_visualization(scramble_map, grid_size):
-    """Creates a PIL image that visualizes the scrambling map."""
     map_size=(512, 512)
     vis_image = Image.new('RGB', map_size, color='lightgray')
     draw = ImageDraw.Draw(vis_image)
@@ -68,13 +64,11 @@ def create_mapping_visualization(scramble_map, grid_size):
     return vis_image
 
 def pil_to_base64(pil_image):
-    """Converts a PIL Image to a Base64 string for HTML embedding."""
     buffered = io.BytesIO()
     pil_image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 def create_canvas_html(base64_string, width, height):
-    """Creates the HTML and JavaScript to render an image on a canvas."""
     return f"""
     <div style="display: flex; justify-content: center; align-items: center; background: #f0f0f0;">
         <canvas id="unscrambled-canvas" width="{width}" height="{height}" style="max-width: 100%; max-height: 512px; object-fit: contain;"></canvas>
@@ -97,10 +91,11 @@ def create_canvas_html(base64_string, width, height):
 
 def process_and_display(input_image, grid_size, seed):
     """
-    Main orchestrator function. Now saves files and returns paths.
+    Main orchestrator function. Saves PNGs and JSON and returns all file paths.
     """
     if input_image is None:
-        return None, "<div>Please upload an image to begin.</div>", None, None
+        # Return empty placeholders for all outputs
+        return None, "<div>Please upload an image to begin.</div>", None, None, None
 
     # 1. Scramble the image
     scrambled_img, scramble_map = scramble_image(input_image, grid_size, seed)
@@ -113,20 +108,31 @@ def process_and_display(input_image, grid_size, seed):
     # 3. Create map visualization
     map_viz_img = create_mapping_visualization(scramble_map, grid_size)
 
-    # --- NEW: Save files and return paths ---
+    # --- Save all files and get their paths ---
     unique_id = uuid.uuid4()
     
-    # Save the scrambled image as a PNG file
+    # Save the scrambled image PNG
     scrambled_filepath = os.path.join(TEMP_DIR, f"{unique_id}_scrambled.png")
     scrambled_img.save(scrambled_filepath)
     
-    # Save the map visualization as a PNG file
+    # Save the map visualization PNG
     map_viz_filepath = os.path.join(TEMP_DIR, f"{unique_id}_map.png")
     map_viz_img.save(map_viz_filepath)
+
+    # NEW: Create and save the JSON map file
+    map_json_filepath = os.path.join(TEMP_DIR, f"{unique_id}_map.json")
+    map_data = {
+        "gridSize": grid_size,
+        "seed": seed,
+        "width": scrambled_img.width,
+        "height": scrambled_img.height,
+        "scrambleMap": scramble_map.tolist() # Convert numpy array for JSON
+    }
+    with open(map_json_filepath, 'w') as f:
+        json.dump(map_data, f, indent=2)
     
-    # 4. Return the file paths and HTML to the Gradio components
-    # The file path is returned twice: once for the gr.Image display and once for the gr.File download.
-    return scrambled_filepath, canvas_html, map_viz_filepath, scrambled_filepath
+    # 4. Return all necessary file paths and the canvas HTML
+    return scrambled_filepath, canvas_html, map_viz_filepath, scrambled_filepath, map_json_filepath
 
 
 # --- Gradio UI Definition ---
@@ -134,9 +140,9 @@ def process_and_display(input_image, grid_size, seed):
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # 🖼️ Secure Image Scrambler & Viewer (v2)
-        This tool scrambles an image so you can share it securely. The scrambled version is provided as a downloadable `.png` file.
-        You can view the unscrambled original in a special preview window that prevents easy downloading.
+        # 🖼️ Secure Image Scrambler & Viewer (v3)
+        Upload an image to create a scrambled `.png` file and a `.json` map file.
+        The unscrambled original can be viewed in a protected preview window that prevents easy downloading.
         """
     )
     
@@ -149,48 +155,53 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             )
             
             with gr.Accordion("Settings", open=True):
-                grid_size_slider = gr.Slider(
-                    minimum=2, maximum=32, value=8, step=1, label="Grid Size (NxN)"
-                )
-                seed_input = gr.Number(
-                    value=lambda: random.randint(0, 99999), label="Scramble Seed"
-                )
+                grid_size_slider = gr.Slider(minimum=2, maximum=32, value=8, step=1, label="Grid Size (NxN)")
+                seed_input = gr.Number(value=lambda: random.randint(0, 99999), label="Scramble Seed")
             
             submit_btn = gr.Button("Scramble & Process", variant="primary")
 
         with gr.Column(scale=2):
             with gr.Tabs():
-                with gr.TabItem("Scrambled Image"):
+                with gr.TabItem("Downloads"):
+                    gr.Markdown("### Scrambled Image")
                     scrambled_output = gr.Image(
                         label="Scrambled Image Preview",
-                        type="filepath", # This component expects a file path
+                        type="filepath",
                         interactive=False,
                     )
-                    # NEW: File download component
-                    downloadable_file = gr.File(
-                        label="Download Scrambled PNG",
+                    downloadable_png = gr.File(
+                        label="Download Scrambled PNG File",
+                        interactive=False
+                    )
+                    gr.Markdown("---")
+                    gr.Markdown("### Scrambling Map")
+                    downloadable_json = gr.File(
+                        label="Download Map JSON File",
                         interactive=False
                     )
+
                 with gr.TabItem("Unscrambled Preview (Protected)"):
                     unscrambled_canvas = gr.HTML(
                         label="Unscrambled Preview (Not directly downloadable)"
                     )
-                with gr.TabItem("Scrambling Map"):
+                    
+                with gr.TabItem("Map Visualization"):
                     mapping_output = gr.Image(
                         label="Mapping Key Visualization",
-                        type="filepath", # This component also expects a file path
+                        type="filepath",
                         interactive=False
                     )
 
-    # Connect the button to the main function, mapping outputs correctly
+    # Connect the button to the main function, mapping all outputs correctly
     submit_btn.click(
         fn=process_and_display, 
         inputs=[input_image, grid_size_slider, seed_input], 
         outputs=[
-            scrambled_output,      # Receives scrambled_filepath
+            scrambled_output,      # Receives scrambled_filepath for display
             unscrambled_canvas,    # Receives canvas_html
             mapping_output,        # Receives map_viz_filepath
-            downloadable_file      # Receives scrambled_filepath
+            downloadable_png,      # Receives scrambled_filepath for download
+            downloadable_json      # Receives map_json_filepath for download
         ]
     )