如何加载自定义 yolo v-7 模型。
这就是我知道如何加载 yolo v-5 模型的方法:
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5/runs/train/exp15/weights/last.pt', force_reload=True)
我在网上看到视频,他们建议使用这个:
!python detect.py --weights runs/train/yolov7x-custom/weights/best.pt --conf 0.5 --img-size 640 --source final_test_v1.mp4
但我希望它像普通模型一样加载,并给我找到对象的边界框坐标。
这就是我在 yolo v-5 中所做的:
from models.experimental import attempt_load
yolov5_weight_file = r'weights/rider_helmet_number_medium.pt' # ... may need full path
model = attempt_load(yolov5_weight_file, map_location=device)
def object_detection(frame):
img = torch.from_numpy(frame)
img = img.permute(2, 0, 1).float().to(device) #convert to required shape based on index
img /= 255.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
pred = model(img, augment=False)[0]
pred = non_max_suppression(pred, conf_set, 0.20) # prediction, conf, iou
# print(pred)
detection_result = []
for i, det in enumerate(pred):
if len(det):
for d in det: # d = (x1, y1, x2, y2, conf, cls)
x1 = int(d[0].item())
y1 = int(d[1].item())
x2 = int(d[2].item())
y2 = int(d[3].item())
conf = round(d[4].item(), 2)
c = int(d[5].item())
detected_name = names[c]
# print(f'Detected: {detected_name} conf: {conf} bbox: x1:{x1} y1:{y1} x2:{x2} y2:{y2}')
detection_result.append([x1, y1, x2, y2, conf, c])
frame = cv2.rectangle(frame, (x1, y1), (x2, y2), (255,0,0), 1) # box
if c!=1: # if it is not head bbox, then write use putText
frame = cv2.putText(frame, f'{names[c]} {str(conf)}', (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1, cv2.LINE_AA)
return (frame, detection_result)
使用torch.hub使用yolov7进行预测
!# Download YOLOv7 code
!git clone https://github.com/WongKinYiu/yolov7
%cd yolov7
from pathlib import Path
import torch
from models.yolo import Model
from utils.general import check_requirements, set_logging
from utils.google_utils import attempt_download
from utils.torch_utils import select_device
dependencies = ['torch', 'yaml']
check_requirements(Path("/content/yolov7/").parent / 'requirements.txt', exclude=('pycocotools', 'thop'))
set_logging()
def custom(path_or_model='path/to/model.pt', autoshape=True):
"""custom mode
Arguments (3 options):
path_or_model (str): 'path/to/model.pt'
path_or_model (dict): torch.load('path/to/model.pt')
path_or_model (nn.Module): torch.load('path/to/model.pt')['model']
Returns:
pytorch model
"""
model = torch.load(path_or_model, map_location=torch.device('cpu')) if isinstance(path_or_model, str) else path_or_model # load checkpoint
if isinstance(model, dict):
model = model['ema' if model.get('ema') else '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
if autoshape:
hub_model = hub_model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available
return hub_model.to(device)
model = custom(path_or_model='yolov7.pt') # custom example
# model = create(name='yolov7', pretrained=True, channels=3, classes=80, autoshape=True) # pretrained example
# Verify inference
import numpy as np
from PIL import Image
imgs = [np.zeros((640, 480, 3))]
results = model(imgs) # batched inference
results.print()
results.save()
df_prediction = results.pandas().xyxy
df_prediction
链接到合作实验室
https://colab.research.google.com/drive/1nKoC-_areXmc_20Xn7z6kcqHEKU7SJsX#scrollTo=yyB_MQW1OWhZ
您可以通过以下方式做到这一点:
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
path = '/path/to/your/file.pt'
model = torch.hub.load("WongKinYiu/yolov7","custom",f"{path}",trust_repo=True)
要获得结果,您可以运行
results = model("/path/to/your/photo")
要获取 bbox,您可以使用:
results.pandas().xyxy
编辑
我使用 python 包创建了一个存储库,以便轻松实现此目的
您不能使用 Yolov5 存储库中的
attempt_load
,因为此方法指向 ultralytics 发布文件。您需要使用 Yolov7 存储库中的 attempt_load
,因为它指向正确的文件。
# yolov7
def attempt_download(file, repo='WongKinYiu/yolov7'):
# Attempt file download if does not exist
file = Path(str(file).strip().replace("'", '').lower())
...
# yolov5
def attempt_download(file, repo='ultralytics/yolov5', release='v6.2'):
# Attempt file download from GitHub release assets if not found locally. release = 'latest', 'v6.2', etc.
from utils.general import LOGGER
def github_assets(repository, version='latest'):
...
然后你可以像这样下载:
# load yolov7 method
from models.experimental import attempt_load
model = attempt_load('yolov7.pt', map_location='cuda:0') # load FP32 model
import torch as th
def loadModel(path:str):
model = th.hub.load("WongKinYiu/yolov7","custom",f{path}",trust_repo=True)
这会起作用。 trust_repo = True 不会要求说出 y 或 n。 在路径中,您可以添加自定义火车模型,例如 ./best.pt
我想使用yolov7-tiny模型从两侧的两个网络摄像头检测水果(已经训练好的新鲜苹果,烂苹果和新鲜番茄,烂番茄)。如果两边的水果都是新鲜的,那么水果就是新鲜的;否则,可以通过双网络摄像头从两侧(即右侧和左侧)检测到腐烂。使用传送带中的伺服电机,我们将新鲜和腐烂的水果分成三个盒子;如果双网络摄像头检测到新鲜苹果,则分开放入第一个盒子中,双网络摄像头检测到新鲜番茄放入第二个盒子中,否则烂苹果和烂番茄分开放入第三个盒子中。我正在使用树莓派4,这是我最后一年的项目,所以你能通过提供代码来帮助我吗?