feat: add image augmentation tool and update demo interface to include new functionality

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from gradio_client import Client, handle_file
-from PIL import Image
+from PIL import Image, ImageFilter
 import numpy as np
 import os
 import time
@@ -172,37 +172,37 @@ register_model_with_metadata(
     architecture="VIT", dataset="TBA"
 )
 
-def postprocess_simple_prediction(result, class_names):
-    scores = {name: 0.0 for name in class_names}
-    fake_prob = result.get("Fake Probability")
-    if fake_prob is not None:
-        # Assume class_names = ["AI", "REAL"]
-        scores["AI"] = float(fake_prob)
-        scores["REAL"] = 1.0 - float(fake_prob)
-    return scores
-
-def simple_prediction(img):
-    client = Client("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview")
-    client.view_api()
-    print(type(img))
-    result = client.predict(
-           handle_file(img),
-           api_name="simple_predict"
-    )
-    return result
-
-
-register_model_with_metadata(
-    model_id="simple_prediction",
-    model=simple_prediction,
-    preprocess=None,
-    postprocess=postprocess_simple_prediction,
-    class_names=["AI", "REAL"], 
-    display_name="Community Forensics",
-    contributor="Jeongsoo Park",
-    model_path="aiwithoutborders-xyz/CommunityForensics-DeepfakeDet-ViT",
-    architecture="ViT", dataset="GOAT"
-)
+# def postprocess_simple_prediction(result, class_names):
+#     scores = {name: 0.0 for name in class_names}
+#     fake_prob = result.get("Fake Probability")
+#     if fake_prob is not None:
+#         # Assume class_names = ["AI", "REAL"]
+#         scores["AI"] = float(fake_prob)
+#         scores["REAL"] = 1.0 - float(fake_prob)
+#     return scores
+
+# def simple_prediction(img):
+#     client = Client("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview")
+#     client.view_api()
+#     print(type(img))
+#     result = client.predict(
+#            handle_file(img),
+#            api_name="simple_predict"
+#     )
+#     return result
+
+
+# register_model_with_metadata(
+#     model_id="simple_prediction",
+#     model=simple_prediction,
+#     preprocess=None,
+#     postprocess=postprocess_simple_prediction,
+#     class_names=["AI", "REAL"], 
+#     display_name="Community Forensics",
+#     contributor="Jeongsoo Park",
+#     model_path="aiwithoutborders-xyz/CommunityForensics-DeepfakeDet-ViT",
+#     architecture="ViT", dataset="GOAT"
+# )
 def infer(image: Image.Image, model_id: str, confidence_threshold: float = 0.75) -> dict:
     """Predict using a specific model.
 
@@ -453,20 +453,20 @@ detection_model_eval_playground = gr.Interface(
 
 community_forensics_preview = gr.Interface(
     fn=lambda: gr.load("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview", src="spaces"),
-    inputs=None,
+    inputs=gr.Image(type="filepath"),
     outputs=gr.HTML(), # or gr.Markdown() if it's just text
     title="Quick and simple prediction by our strongest model.",
     description="No ensemble, no context, no agents, just a quick and simple prediction by our strongest model.",
     api_name="quick_predict"
 )
 
-leaderboard = gr.Interface(
-    fn=lambda: "# AI Generated / Deepfake Detection Models Leaderboard: Soon™",
-    inputs=None,
-    outputs=gr.Markdown(),
-    title="Leaderboard",
-    api_name="leaderboard"
-)
+# leaderboard = gr.Interface(
+#     fn=lambda: "# AI Generated / Deepfake Detection Models Leaderboard: Soon™",
+#     inputs=None,
+#     outputs=gr.Markdown(),
+#     title="Leaderboard",
+#     api_name="leaderboard"
+# )
 
 # simple_predict_interface = gr.Interface(
 #     fn=simple_prediction,
@@ -544,6 +544,37 @@ minmax_processing_interface = gr.Interface(
     api_name="tool_minmax_processing"
 )
 
+def augment_image_interface(img, augment_methods, rotate_degrees, noise_level, sharpen_strength):
+    if img is None:
+        raise gr.Error("No image provided for augmentation. Please upload an image.")
+    
+    # Ensure image is PIL Image and in RGB format
+    if not isinstance(img, Image.Image):
+        try:
+            img = Image.fromarray(img)
+        except Exception as e:
+            raise gr.Error(f"Could not convert input to PIL Image: {e}")
+    if img.mode != 'RGB':
+        img = img.convert('RGB')
+
+    augmented_img, _ = augment_image(img, augment_methods, rotate_degrees, noise_level, sharpen_strength)
+    return augmented_img
+
+augmentation_tool_interface = gr.Interface(
+    fn=augment_image_interface,
+    inputs=[
+        gr.Image(label="Upload Image to Augment", sources=['upload', 'webcam'], type='pil'),
+        gr.CheckboxGroup(["rotate", "add_noise", "sharpen"], label="Augmentation Methods"),
+        gr.Slider(0, 360, value=0, step=1, label="Rotate Degrees", visible=True),
+        gr.Slider(0, 100, value=0, step=1, label="Noise Level", visible=True),
+        gr.Slider(0, 200, value=1, step=1, label="Sharpen Strength", visible=True)
+    ],
+    outputs=gr.Image(label="Augmented Image", type='pil'),
+    title="Image Augmentation Tool",
+    description="Apply various augmentation techniques to your image.",
+    api_name="augment_image"
+)
+
 demo = gr.TabbedInterface(
     [
         detection_model_eval_playground,
@@ -552,7 +583,8 @@ demo = gr.TabbedInterface(
         bit_plane_interface,
         ela_interface,
         gradient_processing_interface,
-        minmax_processing_interface
+        minmax_processing_interface,
+        augmentation_tool_interface
     ],
     [
         "Run Ensemble Prediction",
@@ -561,7 +593,8 @@ demo = gr.TabbedInterface(
         "Bit Plane Values",
         "Error Level Analysis (ELA)",
         "Gradient Processing",
-        "MinMax Processing"
+        "MinMax Processing",
+        "Image Augmentation"
     ],
     title="Deepfake Detection & Forensics Tools",
     theme=None,