def detect_cv2_camera(weightfile): import cv2 model = resnet50() model.load_state_dict(torch.load(weightfile)) print('Loading weights from %s... Done!' % (weightfile)) if use_cuda: model.cuda() cap = cv2.VideoCapture(0) # cap = cv2.VideoCapture("./test.mp4") cap.set(3, 1280) cap.set(4, 720) print("Starting the YOLO loop...") while True: ret, img = cap.read() # 获取当前帧 start = time.time() result_img = plot_boxes_cv2(model, img) finish = time.time() print('Predicted in %f seconds.' % (finish - start)) cv2.imshow('Yolov1 demo', result_img) cv2.waitKey(1)
test_dataset = yoloDataset(root=file_root, list_file='voc2007test.txt', train=False, transform=[transforms.ToTensor()]) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4) print('the dataset has %d images' % (len(train_dataset))) print('the batch_size is %d' % batch_size) # ---------------------网络选择--------------------- use_resnet = True if use_resnet: net = resnet50() else: net = vgg16_bn() if resume: print("loading weight from checkpoints/best.pth") net.load_state_dict(torch.load('checkpoints/best.pth')) else: print('loading pre-trined model ......') if use_resnet: resnet = models.resnet50(pretrained=True) new_state_dict = resnet.state_dict() dd = net.state_dict() for k in new_state_dict.keys(): print(k) if k in dd.keys() and not k.startswith('fc'):
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4) # val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4) print('the train dataset has %d images' % (len(train_dataset))) # print('the test dataset has %d images' % (len(val_dataset))) print('the batch_size is %d' % batch_size) # --------------------- 网络选择 ---------------------- use_resnet = True if use_resnet: net = resnet50() # 自己定义的resnet 网络 else: net = vgg16_bn() """ models 模型表示官方代码 自己定义的是resnet50 """ if resume: print("loading weight from checkpoints/best.pth") net.load_state_dict(torch.load('checkpoints/best.pth')) else: print('loading pre-trained model ......') if use_resnet: # 对应自己网络结构的加载官方预训练模型参数 offiNet = officalModel.resnet50(pretrained=True) # 官方模型 offiNew_state_dict = offiNet.state_dict() dd = net.state_dict()
boxes, cls_indexs, probs = decoder(pred) for i, box in enumerate(boxes): x1 = int(box[0]*w) x2 = int(box[2]*w) y1 = int(box[1]*h) y2 = int(box[3]*h) cls_index = cls_indexs[i] cls_index = int(cls_index) # convert LongTensor to int prob = probs[i] prob = float(prob) result.append([(x1, y1), (x2, y2), VOC_CLASSES[cls_index], image_name, prob]) return result if __name__ == '__main__': model = resnet50() print('load model...') num_epochs=33 learning_rate=0.0001 model.load_state_dict(torch.load('val_best.pth')) #model.load_state_dict(torch.load('best.pth')) model.eval() model.cuda() image_name = 'imgs/person.jpg' image = cv2.imread(image_name) print('predicting...') result = predict_gpu(model, image_name) # 获取预测值 # 绘图 for left_up, right_bottom, class_name, _, prob in result: print("class_name:%s, prob:%f" % (class_name, prob))