22 lines
630 B
Python
22 lines
630 B
Python
|
import torch
|
|||
|
|
import matplotlib.pyplot as plt
|
|||
|
|
|
|||
|
|
def detector(img, model):
|
|||
|
|
"""_summary_
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
img (str or numpy.ndarray): 图片路径或者像素矩阵
|
|||
|
|
model (_type_): 预加载的模型
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
rtn(numpy.ndarray): 渲染后的图片像素点
|
|||
|
|
pred(numpy.ndarray): 检测而出的目标的坐标点、置信度和类别，shape=[n, 6]
|
|||
|
|
"""
|
|||
|
|
result = model(img)
|
|||
|
|
|
|||
|
|
return result.render()[0], result.pred[0].cpu().numpy()
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
# model = torch.hub.load('/home/zhaojh/workspace/git_space/yolov5/', 'yolov5x', source='local', pretrained=True)
|
|||
|
|
pass
|
|||
|
|
|