feat: add Gradio API integration and ONNX preprocessing functions

app.py

@@ -15,7 +15,8 @@ import concurrent.futures
 import ast
 import torch
 
-from utils.utils import softmax, augment_image, preprocess_resize_256, preprocess_resize_224, postprocess_pipeline, postprocess_logits, postprocess_binary_output, to_float_scalar
+from utils.utils import softmax, augment_image, preprocess_resize_256, preprocess_resize_224, postprocess_pipeline, postprocess_logits, postprocess_binary_output, to_float_scalar, infer_gradio_api, preprocess_gradio_api, postprocess_gradio_api
+from utils.onnx_helpers import preprocess_onnx_input, postprocess_onnx_output
 from forensics.gradient import gradient_processing
 from forensics.minmax import minmax_process
 from forensics.ela import ELA
@@ -90,48 +91,6 @@ CLASS_NAMES = {
 }
 
 
-def infer_gradio_api(image_path):
-    client = Client("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview")
-    result_dict = client.predict(
-        input_image=handle_file(image_path),
-        api_name="/simple_predict"
-    )
-    logger.info(f"Debug: Raw result_dict from Gradio API (model_8): {result_dict}, type: {type(result_dict)}")
-    # result_dict is already a dictionary, no need for ast.literal_eval
-    fake_probability = result_dict.get('Fake Probability', 0.0)
-    logger.info(f"Debug: Parsed result_dict: {result_dict}, Extracted fake_probability: {fake_probability}")
-    return {"probabilities": np.array([fake_probability])} # Return as a numpy array with one element
-
-# New preprocess function for Gradio API
-def preprocess_gradio_api(image: Image.Image):
-    # The Gradio API expects a file path, so we need to save the PIL Image to a temporary file.
-    temp_file_path = "./temp_gradio_input.png"
-    image.save(temp_file_path)
-    return temp_file_path
-
-# New postprocess function for Gradio API (adapting postprocess_binary_output)
-def postprocess_gradio_api(gradio_output, class_names):
-    # gradio_output is expected to be a dictionary like {"probabilities": np.array([fake_prob])}
-    probabilities_array = None
-    if isinstance(gradio_output, dict) and "probabilities" in gradio_output:
-        probabilities_array = gradio_output["probabilities"]
-    elif isinstance(gradio_output, np.ndarray):
-        probabilities_array = gradio_output
-    else:
-        logger.warning(f"Unexpected output type for Gradio API post-processing: {type(gradio_output)}. Expected dict with 'probabilities' or numpy.ndarray.")
-        return {class_names[0]: 0.0, class_names[1]: 1.0}
-
-    logger.info(f"Debug: Probabilities array entering postprocess_gradio_api: {probabilities_array}, type: {type(probabilities_array)}, shape: {probabilities_array.shape}")
-
-    if probabilities_array is None or probabilities_array.size == 0:
-        logger.warning("Probabilities array is None or empty after extracting from Gradio API output. Returning default scores.")
-        return {class_names[0]: 0.0, class_names[1]: 1.0}
-
-    # It should always be a single element array for fake probability
-    fake_prob = float(probabilities_array.item())
-    real_prob = 1.0 - fake_prob
-
-    return {class_names[0]: fake_prob, class_names[1]: real_prob}
 
 def register_model_with_metadata(model_id, model, preprocess, postprocess, class_names, display_name, contributor, model_path, architecture=None, dataset=None):
     entry = ModelEntry(model, preprocess, postprocess, class_names, display_name=display_name, contributor=contributor, model_path=model_path, architecture=architecture, dataset=dataset)
@@ -178,27 +137,6 @@ def get_onnx_model_from_cache(hf_model_id):
         _onnx_model_cache[hf_model_id] = load_onnx_model_and_preprocessor(hf_model_id)
     return _onnx_model_cache[hf_model_id]
 
-def preprocess_onnx_input(image: Image.Image, preprocessor_config: dict):
-    # Preprocess image for ONNX model based on preprocessor_config
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-    
-    # Get image size and normalization values from preprocessor_config or use defaults
-    # Use 'size' for initial resize and 'crop_size' for center cropping
-    initial_resize_size = preprocessor_config.get('size', {'height': 224, 'width': 224})
-    crop_size = preprocessor_config.get('crop_size', initial_resize_size['height'])
-    mean = preprocessor_config.get('image_mean', [0.485, 0.456, 0.406])
-    std = preprocessor_config.get('image_std', [0.229, 0.224, 0.225])
-
-    transform = transforms.Compose([
-        transforms.Resize((initial_resize_size['height'], initial_resize_size['width'])),
-        transforms.CenterCrop(crop_size), # Apply center crop
-        transforms.ToTensor(),
-        transforms.Normalize(mean=mean, std=std),
-    ])
-    input_tensor = transform(image)
-    # ONNX expects numpy array with batch dimension (1, C, H, W)
-    return input_tensor.unsqueeze(0).cpu().numpy()
 
 def infer_onnx_model(hf_model_id, preprocessed_image_np, model_config: dict):
     try:
@@ -229,34 +167,6 @@ def infer_onnx_model(hf_model_id, preprocessed_image_np, model_config: dict):
         # Return a structure consistent with other model errors
         return {"logits": np.array([]), "probabilities": np.array([])}
 
-def postprocess_onnx_output(onnx_output, model_config):
-    # Get class names from model_config
-    # Prioritize id2label, then check num_classes, otherwise default
-    class_names_map = model_config.get('id2label')
-    if class_names_map:
-        class_names = [class_names_map[k] for k in sorted(class_names_map.keys())]
-    elif model_config.get('num_classes') == 1: # Handle models that output a single value (e.g., probability of 'Fake')
-        class_names = ['Fake', 'Real'] # Assume first class is 'Fake' and second 'Real'
-    else:
-        class_names = {0: 'Fake', 1: 'Real'} # Default to Fake/Real if not found or not 1 class
-        class_names = [class_names[i] for i in sorted(class_names.keys())]
-
-    probabilities = onnx_output.get("probabilities")
-    
-    if probabilities is not None:
-        if model_config.get('num_classes') == 1 and len(probabilities) == 2: # Special handling for single output models
-            # The single output is the probability of the 'Fake' class
-            fake_prob = float(probabilities[0])
-            real_prob = float(probabilities[1])
-            return {class_names[0]: fake_prob, class_names[1]: real_prob}
-        elif len(probabilities) == len(class_names):
-            return {class_names[i]: float(probabilities[i]) for i in range(len(class_names))}
-        else:
-            logger.warning("ONNX post-processing: Probabilities length mismatch with class names.")
-            return {name: 0.0 for name in class_names}
-    else:
-        logger.warning("ONNX post-processing failed: 'probabilities' key not found in output.")
-        return {name: 0.0 for name in class_names}
 
 # Register the ONNX quantized model
 # Dummy entry for ONNX model to be loaded dynamically

utils/onnx_helpers.py

@@ -0,0 +1,45 @@
+import numpy as np
+from torchvision import transforms
+from PIL import Image
+import logging
+
+def preprocess_onnx_input(image, preprocessor_config):
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    initial_resize_size = preprocessor_config.get('size', {'height': 224, 'width': 224})
+    crop_size = preprocessor_config.get('crop_size', initial_resize_size['height'])
+    mean = preprocessor_config.get('image_mean', [0.485, 0.456, 0.406])
+    std = preprocessor_config.get('image_std', [0.229, 0.224, 0.225])
+    transform = transforms.Compose([
+        transforms.Resize((initial_resize_size['height'], initial_resize_size['width'])),
+        transforms.CenterCrop(crop_size),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=mean, std=std),
+    ])
+    input_tensor = transform(image)
+    return input_tensor.unsqueeze(0).cpu().numpy()
+
+def postprocess_onnx_output(onnx_output, model_config):
+    logger = logging.getLogger(__name__)
+    class_names_map = model_config.get('id2label')
+    if class_names_map:
+        class_names = [class_names_map[k] for k in sorted(class_names_map.keys())]
+    elif model_config.get('num_classes') == 1:
+        class_names = ['Fake', 'Real']
+    else:
+        class_names = {0: 'Fake', 1: 'Real'}
+        class_names = [class_names[i] for i in sorted(class_names.keys())]
+    probabilities = onnx_output.get("probabilities")
+    if probabilities is not None:
+        if model_config.get('num_classes') == 1 and len(probabilities) == 2:
+            fake_prob = float(probabilities[0])
+            real_prob = float(probabilities[1])
+            return {class_names[0]: fake_prob, class_names[1]: real_prob}
+        elif len(probabilities) == len(class_names):
+            return {class_names[i]: float(probabilities[i]) for i in range(len(class_names))}
+        else:
+            logger.warning("ONNX post-processing: Probabilities length mismatch with class names.")
+            return {name: 0.0 for name in class_names}
+    else:
+        logger.warning("ONNX post-processing failed: 'probabilities' key not found in output.")
+        return {name: 0.0 for name in class_names}

utils/utils.py

@@ -1,3 +1,46 @@
+def infer_gradio_api(image_path):
+    from gradio_client import Client, handle_file
+    import numpy as np
+    import logging
+    logger = logging.getLogger(__name__)
+    client = Client("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview")
+    result_dict = client.predict(
+        input_image=handle_file(image_path),
+        api_name="/simple_predict"
+    )
+    logger.info(f"Debug: Raw result_dict from Gradio API (model_8): {result_dict}, type: {type(result_dict)}")
+    fake_probability = result_dict.get('Fake Probability', 0.0)
+    logger.info(f"Debug: Parsed result_dict: {result_dict}, Extracted fake_probability: {fake_probability}")
+    return {"probabilities": np.array([fake_probability])}
+
+def preprocess_gradio_api(image):
+    temp_file_path = "./temp_gradio_input.png"
+    image.save(temp_file_path)
+    return temp_file_path
+
+def postprocess_gradio_api(gradio_output, class_names):
+    import numpy as np
+    import logging
+    logger = logging.getLogger(__name__)
+    probabilities_array = None
+    if isinstance(gradio_output, dict) and "probabilities" in gradio_output:
+        probabilities_array = gradio_output["probabilities"]
+    elif isinstance(gradio_output, np.ndarray):
+        probabilities_array = gradio_output
+    else:
+        logger.warning(f"Unexpected output type for Gradio API post-processing: {type(gradio_output)}. Expected dict with 'probabilities' or numpy.ndarray.")
+        return {class_names[0]: 0.0, class_names[1]: 1.0}
+
+    logger.info(f"Debug: Probabilities array entering postprocess_gradio_api: {probabilities_array}, type: {type(probabilities_array)}, shape: {getattr(probabilities_array, 'shape', None)}")
+
+    if probabilities_array is None or probabilities_array.size == 0:
+        logger.warning("Probabilities array is None or empty after extracting from Gradio API output. Returning default scores.")
+        return {class_names[0]: 0.0, class_names[1]: 1.0}
+
+    fake_prob = float(probabilities_array.item())
+    real_prob = 1.0 - fake_prob
+    return {class_names[0]: fake_prob, class_names[1]: real_prob}
+
 def preprocess_resize_256(image):
     if image.mode != 'RGB':
         image = image.convert('RGB')