Update app.py

added simplified generative inference

app.py

@@ -66,16 +66,89 @@ def predict(image):
     
     return f"Predicted class: {predicted_class.item()}", sample_images_urls
 
+
+# Simplified Generative Inference 
+def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0.01):
+    """
+    Perform Generative Perceptual Inference on the input image.
+    :param image: Input image as a PIL image or numpy array.
+    :param mode: Either 'increase confidence' or 'ReverseDiffuse'.
+    :param model: Pretrained PyTorch model.
+    :param n_iterations: Number of inference iterations.
+    :param step_size: Step size for gradient-based updates.
+    :return: Processed image and gradient visualization.
+    """
+    # Preprocess image
+    transform = transforms.Compose([
+        transforms.Resize(224),
+        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)
+
+    optimizer = torch.optim.SGD([image_tensor], lr=step_size)
+
+    for _ in range(n_iterations):
+        optimizer.zero_grad()
+        output = model(image_tensor)
+        probs = torch.nn.functional.softmax(output, dim=1)
+
+        # Define inference loss based on mode
+        if mode == "increase confidence":
+            # Push away from the least likely classes
+            _, least_likely_indices = torch.topk(probs, k=2, largest=False)
+            losses = []
+            for idx in least_likely_indices[0]:
+                target = torch.full((1,), idx, dtype=torch.long, device=output.device)
+                loss = torch.nn.CrossEntropyLoss()(output, target)
+                losses.append(loss)
+            loss = torch.stack(losses).mean()  # Average the losses for the least likely classes
+        elif mode == "ReverseDiffuse":
+            # Push away from noisy versions
+            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'.")
+
+        loss.backward()
+        optimizer.step()
+
+    # 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_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 = Image.fromarray((processed_image * 255).astype(np.uint8))
+
+    return processed_image, grad_image
+
+
 # # Create the Gradio interface
 # iface = gr.Interface(fn=predict, inputs=gr.Image(type="numpy"), outputs="text")  # Updated from gr.inputs.Image to gr.Image
 
-# Create the Gradio interface
+# # Create the Gradio interface
+# iface = gr.Interface(
+#     fn=predict,
+#     inputs=gr.Image(type="numpy"),
+#     outputs=[gr.Textbox(label="Predicted Class"), gr.Gallery(label="Class Samples")],
+#     title="ResNet-50 VGGFace2 Classifier"
+# )
+
 iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="numpy"),
-    outputs=[gr.Textbox(label="Predicted Class"), gr.Gallery(label="Class Samples")],
-    title="ResNet-50 VGGFace2 Classifier"
+    fn=lambda image, mode: simple_generative_inference(image, mode, model),
+    inputs=[
+        gr.Image(type="pil"),  # Input image
+        gr.Radio(["increase confidence", "ReverseDiffuse"], label="GPI Mode")  # Mode selection
+    ],
+    outputs=[
+        gr.Image(label="Processed Image"),  # Processed image
+        gr.Image(label="Gradient Visualization")  # Gradient visualization
+    ],
+    title="Generative Perceptual Inference (GPI)"
 )
 
-# Launch the interface
+
 iface.launch()