我刚开始了解 YOLO v5 PyTorch 版本并且能够构建模型，因此我尝试使用这个训练有素的模型实现一个用于实时预测的烧瓶应用程序。
负载模型和预测类
class Model(object):
def __init__(self, model):
self.device = torch_utils.select_device()
print(self.device)
model = torch.load(model, map_location=self.device)['model']
self.half = False and self.device.type != 'cpu'
print('half = ' + str(self.half))
if self.half:
model.half()
# model = model.to(self.device).eval()
model.cuda()
self.loaded_model = model
def predict(self, img):
global session
# img1 = torch.from_numpy(img).to(self.device)
# img = img1.reshape(1, 3, 640, 640)
img = img.half() if self.half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
print(img.ndimension())
if img.ndimension() == 3:
img = img.unsqueeze(0)
print(self.loaded_model)
img = img.to(self.device)
# img = img.half()
self.preds = self.loaded_model(img, augment=False)[0]
print(self.predict())
return self.preds
用于从相机或视频中读取帧的相机类
model = Model("weights/best.pt")
class Camera(object):
def __init__(self):
# self.video = cv2.VideoCapture('facial_exp.mkv')
self.video = cv2.VideoCapture(0)
def __del__(self):
self.video.release()
def get_frame(self):
_, fr = self.video.read()
loader = transforms.Compose([transforms.ToTensor()])
image = cv2.resize(fr, (640, 640), interpolation=cv2.INTER_AREA)
input_im = image.reshape(1, 640, 640, 3)
pil_im = Image.fromarray(fr)
image = loader(pil_im).float()
# image = Variable(image, requires_grad=True)
image = image.unsqueeze(0)
pred = model.predict(input_im)
pred = model.predict(image)
print(pred)
_, jpeg = cv2.imencode('.jpg', fr)
return jpeg.tobytes()
一些注释行是我尝试过的方式，但始终如下行
self.preds = self.loaded_model(img, augment=False)[0] 抛出以下错误
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same
解决此错误的任何想法或指导谢谢。