my_infer_hub.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. import base64
  2. import json
  3. import re
  4. import sys
  5. import time
  6. import paddlehub as hub
  7. import cv2
  8. from PIL import Image
  9. import logging
  10. import numpy as np
  11. use_model = "mobile"
  12. use_gpu = False
  13. # img_data = "./1.jpg"
  14. img_data = "D:/Project/PaddleOCR-release-2.0/train_data/bidi_data/orgs_data/train/text_304.jpg"
  15. np_images = [cv2.imread(img_data)]
  16. def only_recognize():
  17. ocr = hub.Module(name="chinese_text_detection_db_"+use_model)
  18. results = ocr
  19. def only_detect():
  20. ocr = hub.Module(name="chinese_text_detection_db_"+use_model)
  21. results = ocr.detect_text(
  22. images=np_images, # 图片数据,ndarray.shape 为 [H, W, C],BGR格式;
  23. use_gpu=use_gpu, # 是否使用 GPU;若使用GPU,请先设置CUDA_VISIBLE_DEVICES环境变量
  24. output_dir='../ocr_result', # 图片的保存路径,默认设为 ocr_result;
  25. visualization=True, # 是否将识别结果保存为图片文件;
  26. box_thresh=0.5 # 检测文本框置信度的阈值;
  27. ) # 识别中文文本置信度的阈值;
  28. for result in results:
  29. print(results)
  30. def detect_and_recognize():
  31. ocr = hub.Module(name="chinese_ocr_db_crnn_"+use_model)
  32. results = ocr.recognize_text(
  33. images=np_images, # 图片数据,ndarray.shape 为 [H, W, C],BGR格式;
  34. use_gpu=use_gpu, # 是否使用 GPU;若使用GPU,请先设置CUDA_VISIBLE_DEVICES环境变量
  35. output_dir='../ocr_result', # 图片的保存路径,默认设为 ocr_result;
  36. visualization=True, # 是否将识别结果保存为图片文件;
  37. box_thresh=0.5, # 检测文本框置信度的阈值;
  38. text_thresh=0.0) # 识别中文文本置信度的阈值;
  39. for result in results:
  40. data = result['data']
  41. save_path = result['save_path']
  42. for infomation in data:
  43. if infomation['text'] == "":
  44. print("no text")
  45. continue
  46. print('text: ', infomation['text'], '\nconfidence: ',
  47. infomation['confidence'], '\ntext_box_position: ',
  48. infomation['text_box_position'])
  49. def image_bigger():
  50. img = cv2.imread("./1.jpg", -1)
  51. height, weight = img.shape[:2]
  52. print(height, weight)
  53. fx = 0.7
  54. fy = 1.2
  55. enlarge = cv2.resize(img, (0, 0), fx=fx, fy=fy, interpolation=cv2.INTER_CUBIC)
  56. print(enlarge.shape[:2])
  57. cv2.imwrite("./1_1.jpg", enlarge)
  58. # img.save("./1_1.jpg", "jpeg")
  59. if __name__ == '__main__':
  60. only_detect()
  61. # image_bigger()
  62. # detect_and_recognize()