fix: pil image generation

app.py

@@ -75,10 +75,10 @@ with gr.Blocks() as demo:
 
         # Right column: Generated image
         with gr.Column(scale=1):
-            output_image = gr.Image(label="Generated Image", height=400, width=400)
+            output_image = gr.Image(label="Generated Image", height=308, width=328)
 
     # Handle click events
-    coord_selector.select(process_click, inputs=[image_idx], outputs=output_image)
+    coord_selector.select(process_click, inputs=[image_idx], outputs=output_image, trigger_mode="multiple")
 
 # Launch the app
 demo.launch()

inference.py

@@ -17,13 +17,13 @@ from huggingface_hub import snapshot_download
 # lpn imports
 from src.models.lpn import LPN
 from src.models.transformer import EncoderTransformer, DecoderTransformer
-from src.visualization import display_grid, ax_to_pil
+from src.visualization import display_grid
 
-from utils import patch_target
+from utils import patch_target, ax_to_pil
 
 
 checkpoint_name = "quiet-thunder-789--checkpoint:v0"
-BLUE_LOCATION_INPUTS = {1: 13, 2: 9}
+BLUE_LOCATION_INPUTS = {0: 13, 1: 9}
 
 local_dir = snapshot_download(repo_id="clement-bonnet/lpn-2d", allow_patterns=f"{checkpoint_name}/*")
 with open(f"{local_dir}/{checkpoint_name}/config.yaml", "r") as f:
@@ -93,4 +93,5 @@ def generate_image(image_idx: int, x: float, y: float, eps: float = 1e-4) -> Ima
     output_grid = output_grids[0]
     _, ax = plt.subplots(1, 1, figsize=(4, 4))
     display_grid(ax=ax, grid=output_grid, grid_shape=jnp.array([4, 4]))
-    return ax_to_pil(ax)
+    pil_image = ax_to_pil(ax)
+    return pil_image

inference_test.py

@@ -0,0 +1,31 @@
+import unittest
+
+from PIL import Image
+
+from inference import generate_image
+
+
+class TestGenerateImage(unittest.TestCase):
+    def test_generate_image_output_type(self):
+        img = generate_image(image_idx=0, x=0.5, y=0.5)
+        self.assertIsInstance(img, Image.Image)
+
+    def test_generate_image_valid_coordinates(self):
+        img = generate_image(image_idx=0, x=0.1, y=0.9)
+        self.assertIsInstance(img, Image.Image)
+
+    def test_generate_image_edge_coordinates(self):
+        img = generate_image(image_idx=1, x=0.0, y=1.0)
+        self.assertIsInstance(img, Image.Image)
+
+    def test_generate_image_invalid_image_idx(self):
+        with self.assertRaises(KeyError):
+            generate_image(image_idx=2, x=0.5, y=0.5)
+
+    def test_generate_image_eps_boundary(self):
+        img = generate_image(image_idx=0, x=1e-5, y=1 - 1e-5)
+        self.assertIsInstance(img, Image.Image)
+
+
+if __name__ == "__main__":
+    unittest.main()

utils.py

@@ -1,3 +1,7 @@
+import io
+
+from PIL import Image
+import matplotlib.pyplot as plt
 import omegaconf
 
 
@@ -10,3 +14,21 @@ def patch_target(config):
         elif isinstance(value, str) and value.startswith("src_v2"):
             # Update the value if it matches the old_value
             config[key] = value.replace("src_v2", "src")
+
+
+def ax_to_pil(ax):
+    fig = ax.figure
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+    buf.seek(0)
+
+    # Load the image data completely before closing the buffer
+    pil_image = Image.open(buf)
+    pil_image_copy = pil_image.copy()
+
+    # Now we can safely close everything
+    pil_image.close()
+    buf.close()
+    plt.close(fig)
+
+    return pil_image_copy