How to use the pytorch model? about pytorch-yolov4 HOT 4 CLOSED

you can refer to models.py

if __name__ == "__main__":
    model = Yolov4()

    pretrained_dict = torch.load('weight/yolov4.pth')
    model.load_state_dict(pretrained_dict)

    num_classes = 80
    if num_classes == 20:
        namesfile = 'data/voc.names'
    elif num_classes == 80:
        namesfile = 'data/coco.names'
    else:
        namesfile = 'data/names'

    use_cuda = 0
    if use_cuda:
        model.cuda()

    from PIL import Image

    img = Image.open("data/dog.jpg").convert('RGB')
    sized = img.resize((608, 608))
    from tool.utils import *

    boxes = do_detect(model, sized, 0.5, 0.4, use_cuda)

    class_names = load_class_names(namesfile)
    plot_boxes(img, boxes, 'predictions.jpg', class_names)

from pytorch-yolov4.

If you use the --load argument when calling train.py, then you can just pass the location of your model.pth

from pytorch-yolov4.

Or use @Tianxiaomo's answer if you want to apply the model to an image

from pytorch-yolov4.

you can refer to models.py

if __name__ == "__main__":
    model = Yolov4()

    pretrained_dict = torch.load('weight/yolov4.pth')
    model.load_state_dict(pretrained_dict)

    num_classes = 80
    if num_classes == 20:
        namesfile = 'data/voc.names'
    elif num_classes == 80:
        namesfile = 'data/coco.names'
    else:
        namesfile = 'data/names'

    use_cuda = 0
    if use_cuda:
        model.cuda()

    from PIL import Image

    img = Image.open("data/dog.jpg").convert('RGB')
    sized = img.resize((608, 608))
    from tool.utils import *

    boxes = do_detect(model, sized, 0.5, 0.4, use_cuda)

    class_names = load_class_names(namesfile)
    plot_boxes(img, boxes, 'predictions.jpg', class_names)

My bad, I only looked into 'tool', 'models.py' works. Thank you!

from pytorch-yolov4.