41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
import sys
|
|
sys.path.append('/home/thebears/Seafile/Designs/ML')
|
|
from model import Model
|
|
import torch
|
|
device = 'cpu'
|
|
|
|
model_rt_path = '/home/thebears/Seafile/Designs/ML/inaturalist_models/models/'#0210701_202822.json
|
|
newest_model = os.path.join(model_rt_path, max(os.listdir(model_rt_path)).replace('.pth',''))
|
|
with open(newest_model + '.json','r') as nmj:
|
|
model_json = json.load(nmj)
|
|
|
|
cats = model_json['categories']
|
|
cats.sort(key=lambda x: x['new_id'])
|
|
num_cat = len(cats) + 1
|
|
model_type = model_json['model_type']
|
|
model = Model(num_cat, model_type)
|
|
labels = [x['name'] for x in cats]
|
|
model.load_state_dict(
|
|
torch.load(newest_model + '.pth', map_location = torch.device(device))
|
|
)
|
|
model.eval()
|
|
|
|
|
|
|
|
# %%
|
|
|
|
onnx_model_path = "models"
|
|
onnx_model_name = "hbirds.onnx"
|
|
os.makedirs(onnx_model_path, exist_ok=True)
|
|
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
|
|
|
|
|
|
# model export into ONNX format
|
|
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
|
|
torch.onnx.export(model, x, full_model_path, opset_version = 12)
|
|
# %%
|
|
|
|
import cv2
|
|
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
|
|
print("OpenCV model was successfully read. Layer IDs: \n", opencv_net.getLayerNames())
|