def start_video(self, model): camera = cv2.VideoCapture(0) while True: frame = camera.read()[1] if frame is None: continue image_array = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) image_array = cv2.resize(image_array, (300, 300)) image_array = substract_mean(image_array) image_array = np.expand_dims(image_array, 0) predictions = model.predict(image_array) detections = detect(predictions, self.prior_boxes) plot_detections(detections, frame, 0.6, self.arg_to_class, self.colors) cv2.imshow('webcam', frame) if cv2.waitKey(1) & 0xFF == ord('q'): break camera.release() cv2.destroyAllWindows()
def test_net(save_folder, net, cuda, dataset, transform, top_k, im_size=300, thresh=0.05): """Test a Fast R-CNN network on an image database.""" num_images = len(dataset) # all detections are collected into: # all_boxes[cls][image] = N x 5 array of detections in # (x1, y1, x2, y2, score) all_boxes = [[[] for _ in range(num_images)] for _ in range(len(labelmap) + 1)] # timers _t = {'im_detect': Timer(), 'misc': Timer()} output_dir = get_output_dir('ssd300_120000', set_type) det_file = os.path.join(output_dir, 'detections.pkl') # detect = Detect(21, 0, 200, 0.01, .45) for i in range(num_images): im, gt, h, w = dataset.pull_item(i) keras_image = np.squeeze(im) # keras_image = substract_mean(x) keras_image_input = np.expand_dims(keras_image, axis=0) keras_output = net.predict(keras_image_input) detections = detect(keras_output, prior_boxes) # detections = detect.forward(keras_output, prior_boxes) _t['im_detect'].tic() # detections = net(x).data detect_time = _t['im_detect'].toc(average=False) # skip j = 0, because it's the background class detection_size = 21 # for j in range(1, detections.size(1)): for j in range(1, detection_size): dets = detections[0, j, :] mask = np.squeeze(dets[:, 0] > 0.01) dets = dets[mask] if len(dets) == 0: continue boxes = dets[:, 1:] boxes[:, 0] *= w boxes[:, 2] *= w boxes[:, 1] *= h boxes[:, 3] *= h # scores = dets[:, 0].cpu().numpy() scores = dets[:, 0] cls_dets = np.hstack( (boxes, scores[:, np.newaxis])).astype(np.float32, copy=False) all_boxes[j][i] = cls_dets print('im_detect: {:d}/{:d} {:.3f}s'.format(i + 1, num_images, detect_time)) with open(det_file, 'wb') as f: pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL) print('Evaluating detections') evaluate_detections(all_boxes, output_dir, dataset)