fix

app.py

@@ -212,7 +212,7 @@ def process_example(image_path_bef, image_path_aft):
     image_aft_grid, _, _ = add_grid_to_image(image_path_aft, 14)
     return image_bef_grid, image_aft_grid  # Reset selected cells and store original image
 
-def display_image(image_path):
+def get_image_size(image_path):
     w, h = Image.open(image_path).convert('RGB').size
     return w, h
     
@@ -246,14 +246,21 @@ with gr.Blocks() as demo:
 
         with gr.Column(scale=1):
 
-            html_text = f"""
+            html_text_bef = f"""
             <div id="container">
-                <canvas id="canvas" width="{width}" height="{height}"></canvas><img id="canvas-background" style="display:none;"/>
+                <canvas id="before" width="{width}" height="{height}"></canvas><img id="canvas-background" style="display:none;"/>
             </div>
             """
 
-            html_bef = gr.HTML(html_text)
-            html_aft = gr.HTML(html_text)
+            html_text_aft = f"""
+            <div id="container">
+                <canvas id="after" width="{width}" height="{height}"></canvas><img id="canvas-background" style="display:none;"/>
+            </div>
+            """
+
+            html_bef = gr.HTML(html_text_bef)
+            gr.HTML("<br>")
+            html_aft = gr.HTML(html_text_aft)
 
 
             # image_display_with_grid_bef = gr.Image(type="pil", label="Before Image with Grid")
@@ -279,9 +286,9 @@ with gr.Blocks() as demo:
                 examples=[["data/images/CLEVR_default_000572.png", "data/images/CLEVR_semantic_000572.png"], 
                             ["data/images/CLEVR_default_003339.png", "data/images/CLEVR_semantic_003339.png"]],
                 inputs=[image_bef, image_aft],
-                # outputs=[image_display_with_grid_bef, image_display_with_grid_aft],
+                outputs=[width, height],
                 label="Example Images",
-                # fn=process_example,
+                fn=get_image_size,
                 examples_per_page=5
                 )
             
@@ -316,14 +323,14 @@ with gr.Blocks() as demo:
         fn=None,
         inputs=[image_bef],
         outputs=[],
-        _js="(image) => { initializeEditor(); importBackground(image); return []; }",
+        _js="(image) => { initializeEditorBefore(); importBackground(image); return []; }",
     )
 
     image_aft.change(
         fn=None,
         inputs=[image_aft],
         outputs=[],
-        _js="(image) => { initializeEditor(); importBackground(image); return []; }",
+        _js="(image) => { initializeEditorAfter(); importBackground(image); return []; }",
     )
 
 

js/interactive_grid.js

@@ -68,7 +68,7 @@ function drawGrid() {
 }
 
 
-function initializeEditor() {
+function initializeEditorBefore() {
     console.log("initializeEditor");
 
     if (isInitialized) {
@@ -77,7 +77,54 @@ function initializeEditor() {
     isInitialized = true;
 
     image = document.getElementById('image');
-    canvas = document.getElementById('canvas');
+    canvas = document.getElementById('before');
+    ctx = canvas.getContext('2d');
+
+    // Add click event listener to canvas
+    canvas.addEventListener('mousedown', handleMouseDown);
+    canvas.addEventListener('mousemove', handleMouseMove);
+    canvas.addEventListener('mouseup', handleMouseUp);
+    canvas.addEventListener('mouseleave', handleMouseLeave);
+
+    cellSize = canvas.width / gridSize;
+
+    canvas.addEventListener('click', (event) => {
+        const rect = canvas.getBoundingClientRect();
+        const scaleX = canvas.width / rect.width;
+        const scaleY = canvas.height / rect.height;
+        const x = (event.clientX - rect.left) * scaleX;
+        const y = (event.clientY - rect.top) * scaleY;
+        const row = Math.floor(y / cellSize);
+        const col = Math.floor(x / cellSize);
+
+        // If the cell is already selected, it's always allowed to deselect it
+        if (grid[row][col]) {
+            grid[row][col] = false;
+            selectedCells--; // Decrement the selected cell count
+        } else {
+            // Only select a new cell if less than 50 cells are already selected
+            if (selectedCells < 50) {
+                grid[row][col] = true;
+                selectedCells++; // Increment the selected cell count
+            }
+        }
+        drawGrid();
+    });
+
+    drawGrid();
+}
+
+
+function initializeEditorAfter() {
+    console.log("initializeEditor");
+
+    if (isInitialized) {
+        return;
+    }
+    isInitialized = true;
+
+    image = document.getElementById('image');
+    canvas = document.getElementById('after');
     ctx = canvas.getContext('2d');
 
     // Add click event listener to canvas