Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update export_onnx.py #272

Merged
merged 1 commit into from
Apr 22, 2024
+20 −10
Merged

Update export_onnx.py #272

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 20 additions & 10 deletions rtdetr_pytorch/tools/export_onnx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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]):
Expand All @@ -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__':
Expand Down