如何解决如何从张量中获得最高置信度的类?
我想以最高的信心获得课程。 下面是执行分类任务的代码:
names = ['class A','class B','class C']
def classify_face(image):
device = torch.device("cpu")
img = process_image(image)
print('Image processed')
# img = image.unsqueeze_(0)
# img = image.float()
pred = model(img)[0]
# Apply NMS
pred = non_max_suppression(pred,0.4,0.5,classes = [0,1,2],agnostic = None )
if classify:
pred = apply_classifier(pred,modelc,img,im0s)
#print(pred)
model.eval()
model.cpu()
print(pred)
# output = non_max_suppression(output,classes = class_names,agnostic = False)
#_,predicted = torch.max(output[0],1)
#print(predicted.data[0],"predicted")
classification = torch.cat(pred)[:,-1]
index = int(classification)
print(names[index])
return names[index]
当 pred 具有一维张量时,上述代码完美运行。但是如果张量大小超过这个值,我就会收到错误消息。
有 5 个元素:x1
、y1
、x2
、y2
、confidence
和 class
。
例如:
pred = [torch.tensor([[212.38568,117.47020,339.35773,266.00513,0.74144,2.00000],[214.60651,118.50694,339.90192,265.91696,0.94277,0.00000]])]
错误:
Traceback (most recent call last):
File "WEBCAM_DETECT.py",line 172,in <module>
label = classify_face(frame)
File "WEBCAM_DETECT.py",line 154,in classify_face
index = int(classification)
ValueError: only one element tensors can be converted to Python scalars
所以我想以最高的信心参加课程。请让我知道如何做到这一点,或者是否有更好的方法。
解决方法
您似乎总是会关注张量 pred[0]
。所以让 pred
成为:
pred = torch.tensor([[212.38568,117.47020,339.35773,266.00513,0.74144,2.00000],[214.60651,118.50694,339.90192,265.91696,0.94277,0.00000]])
置信度最高的预测指标为:
i = torch.argmax(pred[:,4])
因此,您只需获取该索引处的最后一个值:
pred[i,-1]
类名将是 names[int(pred[i,-1])]
。
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。