Update app.py

app.py

@@ -8,7 +8,7 @@ def split_image_grid(image, grid_size):
     # Convert the image to a NumPy array
     img = np.array(image)
     width, height = img.shape[1], img.shape[0]
-    grid_width, grid_height = grid_size, grid_size
+    grid_width, grid_height = grid_size
 
     # Calculate the size of each grid cell
     cell_width = width // grid_width
@@ -20,49 +20,42 @@ def split_image_grid(image, grid_size):
         for j in range(grid_width):
             left = j * cell_width
             upper = i * cell_height
-            right = (j + 1) * cell_width
-            lower = (i + 1) * cell_height
+            right = left + cell_width
+            lower = upper + cell_height
             frame = img[upper:lower, left:right]
             frames.append(frame)
 
     return frames
 
-def create_zip_file(frames):
-    # Create an in-memory zip file
+def zip_images(images):
+    # Create a BytesIO object to hold the zip file
     zip_buffer = BytesIO()
     with zipfile.ZipFile(zip_buffer, 'w') as zipf:
-        for idx, frame in enumerate(frames):
-            frame_byte_array = BytesIO()
-            # Save the frame as a PNG file in the byte array
-            frame_img = Image.fromarray(frame)
-            frame_img.save(frame_byte_array, format="PNG")
-            zipf.writestr(f"frame_{idx}.png", frame_byte_array.getvalue())
+        for idx, img in enumerate(images):
+            # Save each image to the zip file
+            img_buffer = BytesIO()
+            img = Image.fromarray(img)
+            img.save(img_buffer, format='PNG')
+            img_buffer.seek(0)
+            zipf.writestr(f'image_{idx}.png', img_buffer.getvalue())
 
     zip_buffer.seek(0)
     return zip_buffer
 
-def create_gif(frames):
-    # Create a GIF from the frames
-    gif_buffer = BytesIO()
-    frames_pil = [Image.fromarray(frame) for frame in frames]
-    frames_pil[0].save(gif_buffer, format="GIF", save_all=True, append_images=frames_pil[1:], loop=0)
-    gif_buffer.seek(0)
-    return gif_buffer
-
-def process_image(image, grid_size):
-    frames = split_image_grid(image, grid_size)
-    zip_file = create_zip_file(frames)
-    gif_file = create_gif(frames)
-    return zip_file, gif_file
-
 with gr.Blocks() as demo:
     with gr.Row():
-        image_input = gr.Image(type="pil", label="Input Image")
-        grid_size_slider = gr.Slider(1, 10, value=2, step=1, label="Grid Size")
-    with gr.Row():
-        zip_output = gr.File(label="Download ZIP")
-        gif_output = gr.Image(label="Generated GIF")
-    process_button = gr.Button("Process Image")
-    process_button.click(process_image, inputs=[image_input, grid_size_slider], outputs=[zip_output, gif_output])
+        image_input = gr.Image(label="Input Image")
+        grid_size_input = gr.Slider(1, 10, value=2, step=1, label="Grid Size")
+    split_button = gr.Button("Split Image")
+    zip_output = gr.File(label="Download Zipped Images")
+
+    def process_image(image, grid_size):
+        # Split the image into a grid of frames
+        frames = split_image_grid(image, (grid_size, grid_size))
+        # Zip the frames into a single file
+        zip_file = zip_images(frames)
+        return zip_file
+
+    split_button.click(process_image, inputs=[image_input, grid_size_input], outputs=zip_output)
 
 demo.launch(show_error=True)