Update app.py

app.py

@@ -1,142 +1,197 @@
 import gradio as gr
 import numpy as np
-from PIL import Image
-import json
-import os
+from PIL import Image, ImageDraw, ImageFont
+import base64
+import io
 import random
-import math
+import os
+import uuid
+
+# --- Configuration & Setup ---
+# Create a directory to store temporary output files
+TEMP_DIR = "temp_outputs"
+os.makedirs(TEMP_DIR, exist_ok=True)
+
+
+# --- Core Image Processing Logic (Unchanged) ---
+# ... (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
+    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.")
+    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))
+    scramble_map = rng.permutation(len(tiles))
+    scrambled_image = Image.new('RGB', cropped_image.size)
+    for new_pos_index, original_tile_index in enumerate(scramble_map):
+        i, j = divmod(new_pos_index, grid_size)
+        scrambled_image.paste(tiles[original_tile_index], (j * tile_w, i * tile_h))
+    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)
+    for new_pos_index, original_pos_index in enumerate(scramble_map):
+        i, j = divmod(original_pos_index, grid_size)
+        unscrambled_image.paste(scrambled_tiles[new_pos_index], (j * tile_w, i * tile_h))
+    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)
+    try: font = ImageFont.truetype("arial.ttf", size=max(10, 32 - grid_size * 2))
+    except IOError: font = ImageFont.load_default()
+    tile_w, tile_h = map_size[0] // grid_size, map_size[1] // grid_size
+    for new_pos_index, original_pos_index in enumerate(scramble_map):
+        i, j = divmod(new_pos_index, grid_size)
+        x0, y0 = j * tile_w, i * tile_h
+        draw.rectangle([x0, y0, x0 + tile_w, y0 + tile_h], outline='black')
+        text = str(original_pos_index)
+        text_bbox = draw.textbbox((0, 0), text, font=font)
+        text_w, text_h = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
+        draw.text((x0 + (tile_w - text_w) / 2, y0 + (tile_h - text_h) / 2), text, fill='black', font=font)
+    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>
+    </div>
+    <script>
+        function drawImageOnCanvas() {{
+            const canvas = document.getElementById('unscrambled-canvas');
+            if (!canvas) return;
+            const ctx = canvas.getContext('2d');
+            const img = new Image();
+            img.src = "{base64_string}";
+            img.onload = () => {{ ctx.clearRect(0, 0, canvas.width, canvas.height); ctx.drawImage(img, 0, 0, canvas.width, canvas.height); }};
+            canvas.oncontextmenu = (e) => {{ e.preventDefault(); return false; }};
+        }}
+        setTimeout(drawImageOnCanvas, 100);
+    </script>
+    """
 
-# Create a directory for temporary files if it doesn't exist
-if not os.path.exists("temp"):
-    os.makedirs("temp")
+# --- Main Gradio Function ---
 
-def scramble_image(image_pil, block_size):
+def process_and_display(input_image, grid_size, seed):
     """
-    Scrambles an image by dividing it into blocks and shuffling their positions.
+    Main orchestrator function. Now saves files and returns paths.
+    """
+    if input_image is None:
+        return None, "<div>Please upload an image to begin.</div>", None, None
 
-    Args:
-        image_pil (PIL.Image.Image): The input image.
-        block_size (int): The size of the square blocks to divide the image into.
+    # 1. Scramble the image
+    scrambled_img, scramble_map = scramble_image(input_image, grid_size, seed)
+    
+    # 2. Unscramble for canvas preview
+    unscrambled_img = unscramble_image(scrambled_img, scramble_map, grid_size)
+    base64_unscrambled = pil_to_base64(unscrambled_img)
+    canvas_html = create_canvas_html(base64_unscrambled, unscrambled_img.width, unscrambled_img.height)
 
-    Returns:
-        tuple: A tuple containing:
-            - PIL.Image.Image: The scrambled image.
-            - dict: The JSON data representing the scramble map.
-            - str: The file path to the saved JSON file for download.
-    """
-    if image_pil is None:
-        raise gr.Error("Please upload an image first!")
-
-    # Convert PIL Image to a NumPy array
-    img_np = np.array(image_pil)
-    h, w, _ = img_np.shape
-
-    # Calculate the number of blocks, ignoring edges that don't fit a full block
-    num_blocks_h = h // block_size
-    num_blocks_w = w // block_size
-
-    if num_blocks_h == 0 or num_blocks_w == 0:
-        raise gr.Error(f"Image is too small for a block size of {block_size}. Please use a smaller block size.")
-
-    # Create a list of block indices
-    block_indices = list(range(num_blocks_h * num_blocks_w))
-    # Shuffle the indices to create the scramble map
-    scrambled_indices = random.sample(block_indices, len(block_indices))
-
-    # Create a new blank image array to build the scrambled image
-    scrambled_img_np = np.zeros_like(img_np)
-
-    # Reassemble the image using the scrambled map
-    for i, original_idx in enumerate(scrambled_indices):
-        # Position in the new (scrambled) grid
-        new_row, new_col = divmod(i, num_blocks_w)
-        # Position in the original grid
-        original_row, original_col = divmod(original_idx, num_blocks_w)
-
-        # Get the pixel block from the original image
-        block = img_np[
-            original_row * block_size : (original_row + 1) * block_size,
-            original_col * block_size : (original_col + 1) * block_size
-        ]
-        
-        # Place the block into the new scrambled image
-        scrambled_img_np[
-            new_row * block_size : (new_row + 1) * block_size,
-            new_col * block_size : (new_col + 1) * block_size
-        ] = block
-
-    # Convert the scrambled NumPy array back to a PIL Image
-    scrambled_image_pil = Image.fromarray(scrambled_img_np)
-
-    # --- JSON and File Output ---
-    json_data = {
-        "original_dimensions": {"width": w, "height": h},
-        "block_size": block_size,
-        "scramble_map": scrambled_indices
-    }
-
-    # Save the JSON data to a temporary file
-    json_filepath = os.path.join("temp", "scramble_key.json")
-    with open(json_filepath, 'w') as f:
-        json.dump(json_data, f, indent=4)
-
-    return scrambled_image_pil, json_data, json_filepath
-
-# --- Gradio UI with Styling ---
-# Use a pre-built theme for a modern look
-theme = gr.themes.Soft(
-    primary_hue="blue",
-    secondary_hue="sky"
-).set(
-    # Custom CSS for the file drop zone
-    block_label_background_fill="*primary_50",
-    block_label_border_width="1px",
-    block_label_border_color="*primary_200",
-    block_label_text_color="*primary_500",
-)
-
-# Use gr.Blocks for custom layouts
-with gr.Blocks(theme=theme, css=".gradio-container {max-width: 960px !important}") as demo:
+    # 3. Create map visualization
+    map_viz_img = create_mapping_visualization(scramble_map, grid_size)
+
+    # --- NEW: Save files and return paths ---
+    unique_id = uuid.uuid4()
+    
+    # Save the scrambled image as a PNG file
+    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
+    map_viz_filepath = os.path.join(TEMP_DIR, f"{unique_id}_map.png")
+    map_viz_img.save(map_viz_filepath)
+    
+    # 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
+
+
+# --- Gradio UI Definition ---
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # 🖼️ Image Scrambler & Key Generator
-        Drag and drop an image below, choose a block size, and hit "Scramble" to generate a scrambled version 
-        and a downloadable JSON key file that describes how to unscramble it.
+        # 🖼️ 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.
         """
     )
     
-    with gr.Row(variant="panel"):
-        # Column 1: Inputs
-        with gr.Column(scale=1):
-            # The gr.Image component automatically creates a file drop area
-            image_input = gr.Image(
+    with gr.Row():
+        with gr.Column(scale=1, min_width=350):
+            input_image = gr.Image(
                 type="pil", 
-                label="Drop Your Image Here",
-                height=300
-            )
-            block_size_slider = gr.Slider(
-                minimum=8,
-                maximum=128,
-                step=8,
-                value=32,
-                label="Scramble Block Size (pixels)"
+                label="Upload Image or Paste from Clipboard/URL",
+                sources=["upload", "clipboard"]
             )
-            submit_btn = gr.Button("Scramble Image", variant="primary")
+            
+            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"
+                )
+            
+            submit_btn = gr.Button("Scramble & Process", variant="primary")
 
-        # Column 2: Outputs
         with gr.Column(scale=2):
-            scrambled_output_image = gr.Image(label="Scrambled Result", height=400, interactive=False)
-            with gr.Accordion("View & Download Scramble Key", open=False):
-                json_output_display = gr.JSON(label="Scramble Key Data")
-                # The gr.File component provides the download button
-                json_output_file = gr.File(label="Download Key File (.json)", file_count="single")
-
-    # Define the click action: inputs -> function -> outputs
+            with gr.Tabs():
+                with gr.TabItem("Scrambled Image"):
+                    scrambled_output = gr.Image(
+                        label="Scrambled Image Preview",
+                        type="filepath", # This component expects a file path
+                        interactive=False,
+                    )
+                    # NEW: File download component
+                    downloadable_file = gr.File(
+                        label="Download Scrambled PNG",
+                        interactive=False
+                    )
+                with gr.TabItem("Unscrambled Preview (Protected)"):
+                    unscrambled_canvas = gr.HTML(
+                        label="Unscrambled Preview (Not directly downloadable)"
+                    )
+                with gr.TabItem("Scrambling Map"):
+                    mapping_output = gr.Image(
+                        label="Mapping Key Visualization",
+                        type="filepath", # This component also expects a file path
+                        interactive=False
+                    )
+
+    # Connect the button to the main function, mapping outputs correctly
     submit_btn.click(
-        fn=scramble_image,
-        inputs=[image_input, block_size_slider],
-        outputs=[scrambled_output_image, json_output_display, json_output_file]
+        fn=process_and_display, 
+        inputs=[input_image, grid_size_slider, seed_input], 
+        outputs=[
+            scrambled_output,      # Receives scrambled_filepath
+            unscrambled_canvas,    # Receives canvas_html
+            mapping_output,        # Receives map_viz_filepath
+            downloadable_file      # Receives scrambled_filepath
+        ]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()