def detect(img_path, model_path): file_id = utils.get_file_id(img_path) img, img_raw, scale = read_img(img_path) head_detector = Head_Detector_VGG16(ratios=[1], anchor_scales=[2, 4]) dict = torch.load(opt.caffe_pretrain_path) head_detector.load_state_dict(dict['model']) trainer = Head_Detector_Trainer(head_detector).cuda() trainer.load(model_path) img = at.totensor(img) img = img[None, :, :, :] img = img.cuda().float() st = time.time() pred_bboxes_, _ = head_detector.predict(img, scale, mode='evaluate', thresh=THRESH) et = time.time() tt = et - st print("[INFO] Head detection over. Time taken: {:.4f} s".format(tt)) for i in range(pred_bboxes_.shape[0]): ymin, xmin, ymax, xmax = pred_bboxes_[i, :] utils.draw_bounding_box_on_image_array(img_raw, ymin, xmin, ymax, xmax) cv2.imshow('asdf', img_raw) cv2.waitKey()
def detect(img_path, model_path): file_id = utils.get_file_id(img_path) img, img_raw, scale = read_img(img_path) head_detector = Head_Detector_VGG16(ratios=[1], anchor_scales=[2, 4]) trainer = Head_Detector_Trainer(head_detector).cuda() trainer.load(model_path) img = at.totensor(img) img = img[None, :, :, :] img = img.cuda().float() st = time.time() pred_bboxes_, _ = head_detector.predict(img, scale, mode='evaluate', thresh=THRESH) et = time.time() tt = et - st print("[INFO] Head detection over. Time taken: {:.4f} s".format(tt)) for i in range(pred_bboxes_.shape[0]): ymin, xmin, ymax, xmax = pred_bboxes_[i, :] utils.draw_bounding_box_on_image_array(img_raw, ymin, xmin, ymax, xmax) plt.axis('off') plt.imshow(img_raw) if SAVE_FLAG == 1: plt.savefig(os.path.join(opt.test_output_path, file_id + '.png'), bbox_inches='tight', pad_inches=0) else: plt.show()
def detect(img_path, model_path): file_id = utils.get_file_id(img_path) img, img_raw, scale,scale_ = read_img(img_path) print("dim1",img_raw.ndim) head_detector = Head_Detector_VGG16(ratios=[1], anchor_scales=[2,4]) trainer = Head_Detector_Trainer(head_detector).cuda() trainer.load(model_path) img = at.totensor(img) img = img[None, : ,: ,:] img = img.cuda().float() st = time.time() pred_bboxes_, _ = head_detector.predict(img, scale, mode='evaluate', thresh=THRESH) et = time.time() tt = et - st print ("[INFO] Head detection over. Time taken: {:.4f} s".format(tt)) for i in range(pred_bboxes_.shape[0]): print(i) ymin, xmin, ymax, xmax = pred_bboxes_[i,:] print(ymin, xmin, ymax, xmax) image_raw=Image.fromarray(np.uint8(img_raw)) utils.draw_bounding_box_on_image(image_raw,ymin*scale_, xmin*scale_, ymax*scale_, xmax*scale_) img_raw=np.array(image_raw) image_raw=Image.fromarray(np.uint8(img_raw)) if SAVE_FLAG == 1: #image_raw.save('/home/hx/Project/FCHD-Fully-Convolutional-Head-Detector-master/'+file_id+'_1.png') image_raw.save(write_path+'/'+os.path.basename(img_path)) frame_end = cv2.imread(write_path+'/'+os.path.basename(img_path)) cv2.imshow("frame_end",frame_end) key_end = cv2.waitKey(1) & 0xFF else: image_raw.show()