Update app.py

app.py

@@ -65,7 +65,9 @@ from PIL import Image
 import numpy as np
 
 # Simple Generative Inference function
-def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0.01):
+def simple_generative_inference(
+    image, mode, model, n_iterations=10, step_size=0.01, noise_ratio=0.1, eps=0.1
+):
     """
     Perform Generative Perceptual Inference on the input image.
     :param image: Input image as a PIL image.
@@ -73,6 +75,8 @@ def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0
     :param model: Pretrained PyTorch model.
     :param n_iterations: Number of inference iterations.
     :param step_size: Step size for gradient-based updates.
+    :param noise_ratio: Ratio of noise to be added in ReverseDiffuse mode.
+    :param eps: Constraint on perturbation magnitude.
     :return: Processed image and gradient visualization.
     """
     # Preprocess image
@@ -81,51 +85,83 @@ def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # Adjust normalization as needed
     ])
-    image_tensor = transform(image).unsqueeze(0).requires_grad_(True)
+    image_tensor = transform(image).unsqueeze(0)
+    image_tensor.requires_grad_(True)  # Enable gradient computation for the image tensor
 
     optimizer = torch.optim.SGD([image_tensor], lr=step_size)
 
+    # Define least likely classes for "increase confidence" mode
+    if mode == "increase confidence":
+        with torch.no_grad():
+            output = model(image_tensor)
+            probs = torch.nn.functional.softmax(output, dim=1)
+            _, least_likely_classes = torch.topk(probs, k=5, largest=False, dim=1)
+
+    # Create noisy image (only for ReverseDiffuse mode)
+    if mode == "ReverseDiffuse":
+        noisy_image = image_tensor + torch.randn_like(image_tensor) * noise_ratio
+
     for _ in range(n_iterations):
         optimizer.zero_grad()
         output = model(image_tensor)
 
         # Define inference loss based on mode
         if mode == "increase confidence":
-            loss = -torch.nn.functional.cross_entropy(output, output.softmax(dim=1).argmax(dim=1))
+            losses = []
+            for idx in least_likely_classes[0]:  # Iterate over least likely classes
+                target = torch.full((1,), idx, dtype=torch.long, device=output.device)
+                loss = torch.nn.functional.cross_entropy(output, target)
+                losses.append(loss)
+            loss = torch.stack(losses).mean()  # Average loss over least likely classes
         elif mode == "ReverseDiffuse":
-            noisy_image = image_tensor + torch.randn_like(image_tensor) * 0.1
             loss = torch.nn.functional.mse_loss(image_tensor, noisy_image)
         else:
             raise ValueError("Invalid mode selected. Choose 'increase confidence' or 'ReverseDiffuse'.")
 
+        # Compute gradients and update the image
         loss.backward()
-        optimizer.step()
+        grad = image_tensor.grad.data
+        grad_norm = grad.view(grad.shape[0], -1).norm(dim=1, keepdim=True).view_as(image_tensor)
+        scaled_grad = grad / (grad_norm + 1e-10)
+        image_tensor = torch.clamp(
+            image_tensor + step_size * scaled_grad,
+            min=image_tensor - eps,
+            max=image_tensor + eps
+        )
 
     # Generate gradient visualization
-    grad = image_tensor.grad.data.abs().mean(dim=1).squeeze().cpu().numpy()
-    grad_image = (grad - grad.min()) / (grad.max() - grad.min())  # Normalize to [0, 1]
+    grad_visualization = image_tensor.grad.abs().mean(dim=1).squeeze().cpu().numpy()
+    grad_image = (grad_visualization - grad_visualization.min()) / (grad_visualization.max() - grad_visualization.min())
     grad_image = Image.fromarray((grad_image * 255).astype(np.uint8))
 
     # Convert final processed image back to PIL format
     processed_image = image_tensor.detach().squeeze().permute(1, 2, 0).cpu().numpy()
-    processed_image = (processed_image - processed_image.min()) / (processed_image.max() - processed_image.min())  # Normalize
+    processed_image = (processed_image - processed_image.min()) / (processed_image.max() - processed_image.min())
     processed_image = Image.fromarray((processed_image * 255).astype(np.uint8))
 
     return processed_image, grad_image
 
 # Gradio Interface
 iface = gr.Interface(
-    fn=lambda image, mode: simple_generative_inference(image, mode, model),
+    fn=lambda image, mode, step_size, eps, noise_ratio, n_iterations: simple_generative_inference(
+        image, mode, model, step_size=step_size, eps=eps, noise_ratio=noise_ratio, n_iterations=n_iterations
+    ),
     inputs=[
         gr.Image(type="pil", label="Input Image"),  # Input image
-        gr.Radio(["increase confidence", "ReverseDiffuse"], label="Inference Mode")  # Mode selection
+        gr.Radio(["increase confidence", "ReverseDiffuse"], label="Inference Mode"),  # Mode selection
+        gr.Slider(0.001, 1.0, value=0.01, step=0.001, label="Step Size"),  # Step size
+        gr.Slider(0.001, 0.5, value=0.1, step=0.001, label="Epsilon (eps)"),  # Epsilon constraint
+        gr.Slider(0.0, 0.5, value=0.1, step=0.01, label="Noise Ratio"),  # Noise ratio
+        gr.Slider(1, 100, value=10, step=1, label="Number of Iterations"),  # Number of iterations
     ],
     outputs=[
         gr.Image(label="Processed Image"),  # Processed image
         gr.Image(label="Gradient Visualization")  # Gradient visualization
     ],
-    title="Generative Inference"
+    title="Generative Perceptual Inference (GPI)",
+    description="Perform GPI on input images using adjustable parameters such as step size, epsilon, noise ratio, and number of iterations."
 )
 
 
+
 iface.launch()