ONNX models with different IO sizes

configs/linea/linea_hgnetv2_l.py

@@ -6,6 +6,8 @@
 
 output_dir = 'output/linea_hgnetv2_l'
 
+eval_spatial_size = (640, 640)
+
 # backbone
 backbone = 'HGNetv2_B4'
 param_dict_type = backbone.lower()

configs/linea/linea_hgnetv2_m.py

@@ -6,6 +6,8 @@
 
 output_dir = 'output/linea_hgnetv2_m'
 
+eval_spatial_size = (640, 640)
+
 # backbone
 backbone = 'HGNetv2_B2'
 use_lab = True

configs/linea/linea_hgnetv2_n.py

@@ -6,6 +6,8 @@
 
 output_dir = 'output/linea_hgnetv2_n'
 
+eval_spatial_size = (640, 640)
+
 # backbone
 backbone = 'HGNetv2_B0'
 use_lab = True

configs/linea/linea_hgnetv2_s.py

@@ -6,6 +6,8 @@
 
 output_dir = 'output/linea_hgnetv2_s'
 
+eval_spatial_size = (640, 640)
+
 # backbone
 backbone = 'HGNetv2_B1'
 use_lab = True

models/linea/decoder.py

@@ -297,18 +297,18 @@ def forward(self,
             inter_ref_bbox = distance2bbox(ref_points_initial, self.integral(pred_corners, project), self.reg_scale) 
 
             if self.training or layer_id == self.eval_idx:
-            	scores = self.class_embed[layer_id](output)
-            	scores = self.lqe_layers[layer_id](scores, pred_corners)
-            	dec_out_logits.append(scores)
-            	dec_out_bboxes.append(inter_ref_bbox)
+                scores = self.class_embed[layer_id](output)
+                scores = self.lqe_layers[layer_id](scores, pred_corners)
+                dec_out_logits.append(scores)
+                dec_out_bboxes.append(inter_ref_bbox)
 
             pred_corners_undetach = pred_corners
             if self.training:
-            	ref_points_detach = inter_ref_bbox.detach() 
-            	output_detach = output.detach()
+                ref_points_detach = inter_ref_bbox.detach() 
+                output_detach = output.detach()
             else:
-            	ref_points_detach = inter_ref_bbox
-            	output_detach = output
+                ref_points_detach = inter_ref_bbox
+                output_detach = output
 
         return torch.stack(dec_out_bboxes).permute(0, 2, 1, 3), torch.stack(dec_out_logits).permute(0, 2, 1, 3), 
 
@@ -428,6 +428,20 @@ def generate_anchors(self, spatial_shapes):
 
         return output_proposals, output_proposals_valid
 
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                              missing_keys, unexpected_keys, error_msgs):
+        for key in ['output_proposals', 'output_proposals_mask']:
+            full_key = prefix + key
+            if full_key in state_dict:
+                tensor = state_dict[full_key]
+                buf = getattr(self, key, None)
+                if buf is None:
+                    self.register_buffer(key, torch.empty_like(tensor))
+                elif buf.shape != tensor.shape:
+                    self._buffers[key] = torch.empty_like(tensor)
+        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
+                                      missing_keys, unexpected_keys, error_msgs)
+
 
     def forward(self, feats, targets):
         # flatten feature maps 
@@ -445,13 +459,9 @@ def forward(self, feats, targets):
         memory = torch.cat(memory, 1) # bs, \sum{hxw}, c 
 
         # two-stage
-        if self.training:
-            output_proposals, output_proposals_valid = self.generate_anchors(spatial_shapes)
-            output_proposals = output_proposals.to(memory.device).repeat(bs, 1, 1)
-            output_memory = memory.masked_fill(~output_proposals_valid.to(memory.device), float(0))
-        else:
-            output_proposals = self.output_proposals.repeat(bs, 1, 1)
-            output_memory = memory.masked_fill(self.output_proposals_mask, float(0))
+        output_proposals, output_proposals_valid = self.generate_anchors(spatial_shapes)
+        output_proposals = output_proposals.to(memory.device).repeat(bs, 1, 1)
+        output_memory = memory.masked_fill(~output_proposals_valid.to(memory.device), float(0))
 
         output_memory = self.enc_output_norm(self.enc_output(output_memory))
 

tools/deployment/export_onnx.py

@@ -62,8 +62,12 @@ def forward(self, images, orig_target_sizes):
 
     model = Model()
 
-    data = torch.rand(32, 3, 640, 640)
-    size = torch.tensor([[640, 640]])
+    eval_spatial_size = getattr(cfg, 'eval_spatial_size', (640, 640))
+    if eval_spatial_size is None:
+        eval_spatial_size = (640, 640)
+    input_h, input_w = eval_spatial_size
+    data = torch.rand(1, 3, input_h, input_w)
+    size = torch.tensor([[input_h, input_w]])
     _ = model(data, size)
 
     dynamic_axes = {
@@ -84,7 +88,7 @@ def forward(self, images, orig_target_sizes):
         input_names=['images', 'orig_target_sizes'],
         output_names=['lines', 'scores'],
         dynamic_axes=dynamic_axes,
-        opset_version=16,
+        opset_version=18,
         verbose=False,
         do_constant_folding=True,
     )

tools/inference/onnx_inf.py

@@ -2,6 +2,7 @@
 Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
 """
 import os
+import sys
 import cv2
 import glob
 import numpy as np
@@ -10,6 +11,9 @@
 import torchvision.transforms as T
 from PIL import Image, ImageDraw
 
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '../..'))
+from util.slconfig import SLConfig
+
 
 def resize_with_aspect_ratio(image, size, interpolation=Image.BILINEAR):
     """Resizes an image while maintaining aspect ratio and pads it."""
@@ -44,13 +48,13 @@ def draw(images, lines, scores):
     return result_images
 
 
-def process_image(sess, im_pil):
+def process_image(sess, im_pil, input_size):
     w, h = im_pil.size
     orig_size = torch.tensor([w, h])[None]
 
     transforms = T.Compose(
         [
-            T.Resize((640, 640)),
+            T.Resize((input_size[0], input_size[1])),
             T.ToTensor(),
             T.Normalize(mean=[0.538, 0.494, 0.453], std=[0.257, 0.263, 0.273]),
         ]
@@ -69,7 +73,7 @@ def process_image(sess, im_pil):
     print(f"Image processing complete. Result saved as '{OUTPUT_NAME}.jpg'.")
 
 
-def process_video(sess, video_path):
+def process_video(sess, video_path, input_size):
     cap = cv2.VideoCapture(video_path)
 
     # Get video properties
@@ -96,7 +100,7 @@ def process_video(sess, video_path):
 
         transforms = T.Compose(
             [
-                T.Resize((640, 640)),
+                T.Resize((input_size[0], input_size[1])),
                 T.ToTensor(),
                 T.Normalize(mean=[0.538, 0.494, 0.453], std=[0.257, 0.263, 0.273]),
             ]
@@ -128,15 +132,15 @@ def process_video(sess, video_path):
     out.release()
     print(f"Video processing complete. Result saved as '{OUTPUT_NAME}.mp4'.")
 
-def process_file(sess, file_path):
+def process_file(sess, file_path, input_size):
     # Check if the input file is an image or a video
     try:
         # Try to open the input as an image
         im_pil = Image.open(file_path).convert("RGB")
-        process_image(sess, im_pil)
+        process_image(sess, im_pil, input_size)
     except IOError:
         # Not an image, process as video
-        process_video(sess, file_path)
+        process_video(sess, file_path, input_size)
 
 def main(args):
     # Global variable
@@ -147,6 +151,11 @@ def main(args):
     sess = ort.InferenceSession(args.onnx)
     print(f"Using device: {ort.get_device()}")
 
+    cfg = SLConfig.fromfile(args.config)
+    input_size = getattr(cfg, 'eval_spatial_size', (640, 640))
+    if input_size is None:
+        input_size = (640, 640)
+
     input_path = args.input
     thrh = 0.4 if args.thrh is None else args.thrh
 
@@ -160,11 +169,11 @@ def main(args):
         paths = list(glob.iglob(f"{folder_dir}/*.*"))
         for file_path in paths:
             OUTPUT_NAME = file_path.replace(f'{folder_dir}/', f'{output_dir}/').split('.')[0]
-            process_file(sess, file_path)
+            process_file(sess, file_path, input_size)
     else:
         # Process a file
         OUTPUT_NAME = 'onxx_results'
-        process_file(sess, file_path)
+        process_file(sess, file_path, input_size)
 
 if __name__ == "__main__":
     import argparse