Update inference_onnx.py

inference_onnx.py

@@ -31,6 +31,10 @@ class ONNXInference:
         self.device = device
         self.use_openvino = False
 
+        # Hardcoded input size mapping - based on actual model definitions
+        self.input_size = self._get_input_size(model_arch)
+        print(f"Using input size: {self.input_size} for model {model_arch}")
+
         if device == "cuda":
             # Try CUDA first for GPU
             try:
@@ -47,7 +51,7 @@ class ONNXInference:
                 # Check if CUDA is actually being used
                 if self.execution_provider == "CUDAExecutionProvider":
                     print(f"Using ONNX Runtime with {self.execution_provider}")
-                    # Get transform from timm model
+                    # Get transform with correct input size
                     self.transform = self._get_transform()
                     return
                 else:
@@ -76,7 +80,7 @@ class ONNXInference:
             print(f"OpenVINO initialization failed: {e}, falling back to ONNX Runtime CPU")
             self._init_onnx_runtime_cpu(model_path)
 
-        # Get transform from timm model
+        # Get transform with correct input size
         self.transform = self._get_transform()
 
     def _init_onnx_runtime_cpu(self, model_path):
@@ -92,15 +96,24 @@ class ONNXInference:
         self.execution_provider = self.session.get_providers()[0]
         print(f"Using ONNX Runtime with {self.execution_provider}")
 
-        # Get transform from timm model
-        self.transform = self._get_transform()
+    def _get_input_size(self, model_arch):
+        """Get input size based on model architecture - hardcoded to match actual model definitions"""
+        if model_arch == 'lsnet_xl_artist_448':
+            return 448
+        else:
+            # All other artist models use 224
+            return 224
 
     def _get_transform(self):
-        """Create preprocessing transform from timm model config"""
+        """Create preprocessing transform with correct input size"""
+        # Create a dummy model to get the base config
         model = create_model(self.model_arch, pretrained=False)
         model.eval()
-        config = resolve_data_config({}, model=model)
+        
+        # Override the input size with our hardcoded value
+        config = resolve_data_config({'input_size': (3, self.input_size, self.input_size)}, model=model)
         transform = create_transform(**config)
+        print(f"Created ONNX transform with input size: {self.input_size}")
         return transform
 
     def preprocess(self, image):
@@ -115,6 +128,7 @@ class ONNXInference:
         """
         image = image.convert("RGB")
         tensor = self.transform(image)
+        print(f"Preprocessed image to tensor shape: {tensor.shape}")
         return tensor.unsqueeze(0).cpu().numpy()
 
     def predict(self, image, top_k=5, threshold=0.0):
@@ -148,4 +162,4 @@ class ONNXInference:
 def softmax(x):
     """Compute softmax values for a set of scores."""
     e_x = np.exp(x - np.max(x))
-    return e_x / e_x.sum(axis=0)
+    return e_x / e_x.sum(axis=0)