app.py

import torch
import gradio as gr
from torchvision import models, transforms
from PIL import Image
import requests
from huggingface_hub import hf_hub_download

from PIL import Image
import numpy as np
import random




# Load the model checkpoint from Hugging Face
checkpoint_path = hf_hub_download(repo_id="ttoosi/resnet50_robust_face", filename="100_checkpoint.pt")

# Initialize the model
model = models.resnet50()
# change the num_classes to 500
model.fc = torch.nn.Linear(model.fc.in_features, 500)
checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))['model']
# remove the prefix 'module.' from the keys
# remove the prefix 'model.' from the keys that have it
new_state_dict = {k.replace('module.', ''): v for k, v in checkpoint.items()}
new_state_dict = {k.replace('model.', ''): v for k, v in new_state_dict.items()}
new_state_dict = {k.replace('attacker.', ''): v for k, v in new_state_dict.items()}


print(new_state_dict.keys())
print('********************')
model.load_state_dict(new_state_dict, strict=False)  # ignore Unexpected key(s) in state_dict: "normalizer.new_mean", "normalizer.new_std", "normalize.new_mean", "normalize.new_std". 
model.eval()

# Image preprocessing
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), # vggface2
])

# # Function to make predictions
# def predict(image):
#     if isinstance(image, np.ndarray):
#         image = Image.fromarray(image)  # Convert to PIL Image if i
#     image = preprocess(image).unsqueeze(0)  # Add batch dimension
#     with torch.no_grad():
#         output = model(image)  # Perform inference on CPU
#     _, predicted_class = output.max(1)
#     # Fetch 9 random samples from the predicted class
#     class_samples = ds.filter(lambda example: example['label'] == predicted_class.item())

#     sample_images = random.sample(list(class_samples), min(len(class_samples), 9))
    
#     sample_images_urls = [sample['image'] for sample in sample_images]
    
#     return f"Predicted class: {predicted_class.item()}", sample_images_urls


import torch
import torch.nn.functional as F
from torchvision import transforms
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, noise_ratio=0.1, eps=0.1
):
    """
    Perform Generative Perceptual Inference on the input image.
    :param image: Input image as a PIL image.
    :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.
    :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
    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)
    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":
            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":
            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()
        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_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())
    processed_image = Image.fromarray((processed_image * 255).astype(np.uint8))

    return processed_image, grad_image

# Gradio Interface
iface = gr.Interface(
    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.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 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()