Update app.py

app.py

@@ -101,6 +101,23 @@ def gram_matrix(tensor):
     
     return gram
 
+# Function to resize image while maintaining aspect ratio
+def resize_image(image_path, max_size=400):
+    img = Image.open(image_path).convert('RGB')
+    ratio = max_size / max(img.size)
+    new_size = tuple([int(x*ratio) for x in img.size])
+    img = img.resize(new_size, Image.Resampling.LANCZOS)
+    return np.array(img)
+
+def create_grid(images, rows, cols):
+    assert len(images) == rows * cols, "Number of images doesn't match the grid size"
+    w, h = images[0].shape[1], images[0].shape[0]
+    grid = np.zeros((h*rows, w*cols, 3), dtype=np.uint8)
+    for i, img in enumerate(images):
+        r, c = divmod(i, cols)
+        grid[r*h:(r+1)*h, c*w:(c+1)*w] = img
+    return grid
+
 def style_transfer(content_image, style_image, alpha, beta, conv1_1, conv2_1, conv3_1, conv4_1, conv5_1, steps):
     content = load_image(content_image).to(device)
     style = load_image(style_image, shape=content.shape[-2:]).to(device)
@@ -124,6 +141,9 @@ def style_transfer(content_image, style_image, alpha, beta, conv1_1, conv2_1, co
     
     optimizer = optim.Adam([target], lr=0.003)
     
+    intermediate_images = []
+    show_every = steps // 9  # Show 9 intermediate images
+
     for ii in range(1, steps+1):
         target_features = get_features(target, vgg)
         content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2'])**2)
@@ -142,17 +162,15 @@ def style_transfer(content_image, style_image, alpha, beta, conv1_1, conv2_1, co
         optimizer.zero_grad()
         total_loss.backward()
         optimizer.step()
+        
+        if ii % show_every == 0 or ii == steps:
+            intermediate_images.append(im_convert(target))
     
-    return im_convert(target)
-
+    final_image = intermediate_images[-1]
+    intermediate_grid = create_grid(intermediate_images, 3, 3)
+    
+    return final_image, intermediate_grid
 
-# Function to resize image while maintaining aspect ratio
-def resize_image(image_path, max_size=400):
-    img = Image.open(image_path).convert('RGB')
-    ratio = max_size / max(img.size)
-    new_size = tuple([int(x*ratio) for x in img.size])
-    img = img.resize(new_size, Image.Resampling.LANCZOS)
-    return np.array(img)
 
 # Example images
 # examples = [
@@ -208,6 +226,7 @@ with gr.Blocks() as demo:
     steps_slider = gr.Slider(minimum=1, maximum=2000, value=1000, step=100, label="Number of Steps")
     
     run_button = gr.Button("Run Style Transfer")
+    intermediate_output = gr.Image(label="Intermediate Results")
     
     run_button.click(
         style_transfer,
@@ -223,7 +242,7 @@ with gr.Blocks() as demo:
             conv5_1_slider,
             steps_slider
         ],
-        outputs=output_image
+        outputs=[output_image, intermediate_output]
     )
     
     gr.Examples(