Home > Software design >  Deploy pytorch .pth model in a python script
Deploy pytorch .pth model in a python script

Time:11-21

After successfully training my yolact model using a custom dataset I'm happy with the inference results outputted by eval.py using this command from anaconda terminal:

python eval.py --trained_model=./weights/yolact_plus_resnet50_abrasion_39_10000.pth --config=yolact_resnet_abrasion_config --score_threshold=0.8 --top_k=15 --images=./images:output_images

Now I want to run this inference from my own python script instead of using the anaconda terminal. I wanna be able to get the bounding boxes of detections made on webcam frames obtained by this code below. Any idea ?

import cv2

src = cv2.VideoCapture(0)

while True:
    ret, frame = src.read()
    cv2.imshow('frame', frame)
    key = cv2.waitKey(5)
    if key == (27):
        break

The eval.py code is here at Yolact repository https://github.com/dbolya/yolact/blob/master/eval.py

CodePudding user response:

I will just write the pseudocode here for you.

Step 1: Try loading the model using the lines starting from here and ending here

Step 2: Use this function for evaluation. Instead of cv2.imread, you just need to send your frame

Step 3: Follow this function to get the bounding boxes. Especially this line. Just trackback the 't' variable and you will get your bounding boxes.

Hope it helps. Let me know if you need more clarification.

  • Related