helpers.py
2.84 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
from pathlib import Path
from yolo_module.yolov5.models.yolo import Model
from yolo_module.yolov5.utils.general import set_logging, yolov5_in_syspath
from yolo_module.yolov5.utils.google_utils import attempt_download
from yolo_module.yolov5.utils.torch_utils import torch
def load_model(model_path, device=None, autoshape=True, verbose=False):
"""
Creates a specified YOLOv5 model
Arguments:
model_path (str): path of the model
config_path (str): path of the config file
device (str): select device that model will be loaded (cpu, cuda)
pretrained (bool): load pretrained weights into the model
autoshape (bool): make model ready for inference
verbose (bool): if False, yolov5 logs will be silent
Returns:
pytorch model
(Adapted from yolo_module.yolov5.hubconf.create)
"""
# set logging
set_logging(verbose=verbose)
# set device if not given
if not device:
device = "cuda:0" if torch.cuda.is_available() else "cpu"
attempt_download(model_path) # download if not found locally
with yolov5_in_syspath():
model = torch.load(model_path, map_location=torch.device(device))
if isinstance(model, dict):
model = model["model"] # load model
hub_model = Model(model.yaml).to(next(model.parameters()).device) # create
hub_model.load_state_dict(model.float().state_dict()) # load state_dict
hub_model.names = model.names # class names
model = hub_model
if autoshape:
model = model.autoshape()
return model
class YOLOv5:
def __init__(self, model_path, device=None, load_on_init=True):
self.model_path = model_path
self.device = device
if load_on_init:
Path(model_path).parents[0].mkdir(parents=True, exist_ok=True)
self.model = load_model(model_path=model_path, device=device, autoshape=True)
else:
self.model = None
def load_model(self):
"""
Load yolov5 weight.
"""
Path(self.model_path).parents[0].mkdir(parents=True, exist_ok=True)
self.model = load_model(model_path=self.model_path, device=self.device, autoshape=True)
def predict(self, image_list, size=640, augment=False):
"""
Perform yolov5 prediction using loaded model weights.
Returns results as a yolov5.models.common.Detections object.
"""
assert self.model is not None, "before predict, you need to call .load_model()"
results = self.model(imgs=image_list, size=size, augment=augment)
return results
if __name__ == "__main__":
model_path = "yolov5/weights/yolov5s.pt"
device = "cuda"
model = load_model(model_path=model_path, config_path=None, device=device)
from PIL import Image
imgs = [Image.open(x) for x in Path("yolov5/data/images").glob("*.jpg")]
results = model(imgs)