Update export_onnx.py

When saving a picture, restore the picture size.
Add output label and box details.

Author: NeKoooo233

rtdetr_pytorch/tools/export_onnx.py

@@ -82,13 +82,18 @@ def forward(self, images, orig_target_sizes):
 
 
     # import onnxruntime as ort 
-    # from PIL import Image, ImageDraw
+    # from PIL import Image, ImageDraw, ImageFont
     # from torchvision.transforms import ToTensor
+    # from src.data.coco.coco_dataset import mscoco_category2name, mscoco_category2label, mscoco_label2category
 
     # # print(onnx.helper.printable_graph(mm.graph))
 
-    # im = Image.open('./000000014439.jpg').convert('RGB')
-    # im = im.resize((640, 640))
+    # # Load the original image without resizing
+    # original_im = Image.open('./hongkong.jpg').convert('RGB')
+    # original_size = original_im.size
+
+    # # Resize the image for model input
+    # im = original_im.resize((640, 640))
     # im_data = ToTensor()(im)[None]
     # print(im_data.shape)
 
@@ -104,7 +109,7 @@ def forward(self, images, orig_target_sizes):
 
     # labels, boxes, scores = output
 
-    # draw = ImageDraw.Draw(im)
+    # draw = ImageDraw.Draw(original_im)  # Draw on the original image
     # thrh = 0.6
 
     # for i in range(im_data.shape[0]):
@@ -115,12 +120,17 @@ def forward(self, images, orig_target_sizes):
 
     #     print(i, sum(scr > thrh))
 
-    #     for b in box:
-    #         draw.rectangle(list(b), outline='red',)
-    #         draw.text((b[0], b[1]), text=str(lab[i]), fill='blue', )
-
-    # im.save('test.jpg')
-
+    # for b, l in zip(box, lab):
+    #     # Scale the bounding boxes back to the original image size
+    #     b = [coord * original_size[j % 2] / 640 for j, coord in enumerate(b)]
+    #     # Get the category name from the label
+    #     category_name = mscoco_category2name[mscoco_label2category[l]]
+    #     draw.rectangle(list(b), outline='red', width=2)
+    #     font = ImageFont.truetype("Arial.ttf", 15)
+    #     draw.text((b[0], b[1]), text=category_name, fill='yellow', font=font)
+
+    # # Save the original image with bounding boxes
+    # original_im.save('test.jpg')
 
 
 if __name__ == '__main__':