Update app.py

app.py

@@ -9,18 +9,9 @@ from PIL import Image
 import numpy as np
 import random
 
-# from datasets import load_dataset
-# from datasets import DatasetDict
-# ds = DatasetDict({
-#     "validation": load_dataset("chronopt-research/cropped-vggface2-224", split="validation"),
-# })
 
 
 
-# Load the VGGFace2 dataset using Hugging Face's datasets library
-# ds = load_dataset("chronopt-research/cropped-vggface2-224", split="validation")
-
-
 # Load the model checkpoint from Hugging Face
 checkpoint_path = hf_hub_download(repo_id="ttoosi/resnet50_robust_face", filename="100_checkpoint.pt")
 
@@ -67,11 +58,17 @@ preprocess = transforms.Compose([
 #     return f"Predicted class: {predicted_class.item()}", sample_images_urls
 
 
-# Simplified Generative Inference 
+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):
     """
     Perform Generative Perceptual Inference on the input image.
-    :param image: Input image as a PIL image or numpy array.
+    :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.
@@ -91,20 +88,11 @@ def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0
     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
+            loss = -torch.nn.functional.cross_entropy(output, output.softmax(dim=1).argmax(dim=1))
         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:
@@ -125,29 +113,18 @@ def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0
 
     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
-# 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"
-# )
-
+# Gradio Interface
 iface = gr.Interface(
     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
+        gr.Image(type="pil", label="Input Image"),  # Input image
+        gr.Radio(["increase confidence", "ReverseDiffuse"], label="Inference Mode")  # Mode selection
     ],
     outputs=[
         gr.Image(label="Processed Image"),  # Processed image
         gr.Image(label="Gradient Visualization")  # Gradient visualization
     ],
-    title="Generative Perceptual Inference (GPI)"
+    title="Generative Inference"
 )