mihaidusmanu / d2-net

D2-Net: A Trainable CNN for Joint Description and Detection of Local Features
Other
761 stars 163 forks source link

please provide source code to export onnx file #106

Open ChengLiu-xsy opened 6 months ago

ChengLiu-xsy commented 6 months ago

hi: Could you please provide the source code of exporting onnx file?

BADAL244 commented 4 months ago

import torch from lib.model_test import D2Net from lib.utils import preprocess_image from lib.pyramid import process_multiscale

CUDA

use_cuda = torch.cuda.is_available() device = torch.device("cuda:0" if use_cuda else "cpu")

Initialize the model

model_file = '/home/badal/d2-net/d2_tf.pth' # Specify the path to your model file use_relu = True # Specify whether to use ReLU

model = D2Net( model_file=model_file, use_relu=use_relu, use_cuda=use_cuda )

Ensure the model is in evaluation mode

model.eval()

dummy_input = torch.randn(1, 3, 480, 640, device=device)

Specify the path to save the ONNX model

onnx_file_path = '/home/badal/d2-net/onnx/model_new.onnx'

Export the model

torch.onnx.export(model, dummy_input, onnx_file_path, export_params=True, opset_version=11, input_names=['input'], output_names=['descriptor' , 'score' , 'keypoint' ])

print(f"Model has been converted to ONNX and saved to {onnx_file_path}")

BADAL244 commented 4 months ago

https://drive.google.com/file/d/17B-VuCAcNt7Mm40y3rJxEth5Nl6tTF8R/view?usp=sharing