How to load pytorch pth model correctly and combine with another model?

user_8614839 0
Viewed 2093

I am looking for an OCR pretrained model that works with pytorch. I tried https://github.com/clovaai/CRAFT-pytorch, but it is not supported on pytorch hub. And I cannot load pth model, because it has only weights. How to load model?

My first model is a yolov5 model trained on custom data, so it should crop an image and send it to the next model. The next model should be OCR, mostly digit recognition. But I can't run craft-pytorch

model = torch.hub.load('.', 'custom', path='runs/train/exp2/weights/best.pt', source='local', force_reload=True)

# It throws an error with pytorch hub
model_ocr = torch.hub.load('clovaai/CRAFT-pytorch', 'craft_mlt_25k.pth')
# Tried with torch load, but pth files have only weights not model
ocr_model = torch.load('runs/ocr/craft_mlt_25k.pth')

cap = cv2.VideoCapture('../Dataset/test/09-10.mp4')

while(cap.isOpened()):
    ret, frame = cap.read()
    results = model(frame)
    crops = results.crop(save=False)    
    for crop in crops:
        if 'number' in crop['label']:
            ocr_result = model_ocr(crop['im'])
            ocr_crop = ocr_result.crop(save=False)
pytorch ocr python
1 Answers

If you have the source code for the model, you can create it and load its weights, but as I understood the question, you wish to do it without the source code.

For that, you can just call torch.load:

state_dict = torch.load('model_file.pth')

state_dict should just be a (type of) dictionary. You can examine and pick it apart, as you wish:

print(state_dict.keys())

with torch.no_grad():
    my_model.some_weight[:] = state_dict['some_other_name']

(Note that it's not safe to call torch.load on untrusted data)

user_19371970
Related