fix: add missing import for infer_onnx_model in model_loader and update its definition in onnx_helpers

app.py

@@ -16,7 +16,7 @@ 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, infer_gradio_api, preprocess_gradio_api, postprocess_gradio_api
-from utils.onnx_helpers import preprocess_onnx_input, postprocess_onnx_output
+from utils.onnx_helpers import preprocess_onnx_input, postprocess_onnx_output, infer_onnx_model
 from utils.model_loader import register_all_models
 from utils.onnx_model_loader import load_onnx_model_and_preprocessor, get_onnx_model_from_cache
 from forensics.gradient import gradient_processing
@@ -93,34 +93,6 @@ def register_model_with_metadata(model_id, model, preprocess, postprocess, class
 _onnx_model_cache = {}
 
 
-def infer_onnx_model(hf_model_id, preprocessed_image_np, model_config: dict):
-    try:
-        ort_session, _, _ = get_onnx_model_from_cache(hf_model_id, _onnx_model_cache, load_onnx_model_and_preprocessor)
-        
-        # Debug: Print expected input shape from ONNX model
-        for input_meta in ort_session.get_inputs():
-            logger.info(f"Debug: ONNX model expected input name: {input_meta.name}, shape: {input_meta.shape}, type: {input_meta.type}")
-
-        logger.info(f"Debug: preprocessed_image_np shape: {preprocessed_image_np.shape}")
-        ort_inputs = {ort_session.get_inputs()[0].name: preprocessed_image_np}
-        ort_outputs = ort_session.run(None, ort_inputs)
-        
-        logits = ort_outputs[0]
-        logger.info(f"Debug: logits type: {type(logits)}, shape: {logits.shape}")
-        # If the model outputs a single logit (e.g., shape (1,)), use .item() to convert to scalar
-        # Otherwise, assume it's a batch of logits (e.g., shape (1, num_classes)) and take the first element (batch dim)
-        # The num_classes in config.json can be misleading; rely on actual output shape.
-        
-        # Apply softmax to the logits to get probabilities for the classes
-        # The softmax function in utils/utils.py now ensures a list of floats
-        probabilities = softmax(logits[0]) # Assuming logits[0] is the relevant output for a single prediction
-
-        return {"logits": logits, "probabilities": probabilities}
-
-    except Exception as e:
-        logger.error(f"Error during ONNX inference for {hf_model_id}: {e}")
-        # Return a structure consistent with other model errors
-        return {"logits": np.array([]), "probabilities": np.array([])}
 
 # Register all models (ONNX, HuggingFace, Gradio API)
 register_all_models(MODEL_PATHS, CLASS_NAMES, device, infer_onnx_model, preprocess_onnx_input, postprocess_onnx_output)

utils/model_loader.py

@@ -1,5 +1,5 @@
 # Add missing import for infer_onnx_model
-from app import infer_onnx_model
+from utils.onnx_helpers import infer_onnx_model
 # Add missing import for preprocess_onnx_input
 from utils.onnx_helpers import preprocess_onnx_input
 """

utils/onnx_helpers.py

@@ -1,3 +1,25 @@
+# Shared ONNX inference function for use by app.py and model_loader.py
+def infer_onnx_model(hf_model_id, preprocessed_image_np, model_config: dict):
+    from .onnx_model_loader import get_onnx_model_from_cache, load_onnx_model_and_preprocessor
+    from .utils import softmax
+    import numpy as np
+    import logging
+    logger = logging.getLogger(__name__)
+    _onnx_model_cache = {}
+    try:
+        ort_session, _, _ = get_onnx_model_from_cache(hf_model_id, _onnx_model_cache, load_onnx_model_and_preprocessor)
+        for input_meta in ort_session.get_inputs():
+            logger.info(f"Debug: ONNX model expected input name: {input_meta.name}, shape: {input_meta.shape}, type: {input_meta.type}")
+        logger.info(f"Debug: preprocessed_image_np shape: {preprocessed_image_np.shape}")
+        ort_inputs = {ort_session.get_inputs()[0].name: preprocessed_image_np}
+        ort_outputs = ort_session.run(None, ort_inputs)
+        logits = ort_outputs[0]
+        logger.info(f"Debug: logits type: {type(logits)}, shape: {logits.shape}")
+        probabilities = softmax(logits[0])
+        return {"logits": logits, "probabilities": probabilities}
+    except Exception as e:
+        logger.error(f"Error during ONNX inference for {hf_model_id}: {e}")
+        return {"logits": np.array([]), "probabilities": np.array([])}
 import numpy as np
 from torchvision import transforms
 from PIL import Image