test_api.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import requests
  2. from io import BytesIO
  3. import base64
  4. import json
  5. from PIL import Image, ImageFilter, ImageEnhance
  6. from config import PROMPT_EXTRACT_ICON, PROMPT_EXTRACT_NAME, PROMPT_EXTRACT_COMPONENTS, PROMPT_EXTRACT_KEYWORD, PROMPT_EXTRACT_PREVENTION, PROMPT_EXTRACT_SUPPLIER
  7. def image_to_base64(pil_image, image_format="JPEG"):
  8. """将PIL Image图像转换为Base64编码"""
  9. buffered = BytesIO()
  10. pil_image.save(buffered, format=image_format)
  11. img_byte_array = buffered.getvalue()
  12. encode_image = base64.b64encode(img_byte_array).decode('utf-8')
  13. return encode_image
  14. def resize_image(image, max_size=512):
  15. """缩放图像尺寸,保持 OCR 质量"""
  16. width, height = image.size
  17. max_dim = max(width, height)
  18. # 如果图像不需要缩小,直接返回
  19. if max_dim <= max_size:
  20. return image
  21. scaling_factor = max_size / max_dim
  22. new_width = int(width * scaling_factor)
  23. new_height = int(height * scaling_factor)
  24. resized = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
  25. resized = resized.filter(ImageFilter.UnsharpMask(radius=1, percent=120, threshold=3))
  26. enhancer = ImageEnhance.Contrast(resized)
  27. resized = enhancer.enhance(1.1)
  28. return resized
  29. image = Image.open('./test1.jpg')
  30. image = resize_image(image)
  31. image_base64 = image_to_base64(image)
  32. # response = requests.post(
  33. # "http://127.0.0.1:8000/api/v1/ocr",
  34. # json={
  35. # "image": image_base64,
  36. # "text": [PROMPT_EXTRACT_ICON, PROMPT_EXTRACT_NAME, PROMPT_EXTRACT_COMPONENTS, PROMPT_EXTRACT_KEYWORD, PROMPT_EXTRACT_PREVENTION, PROMPT_EXTRACT_SUPPLIER]
  37. # }
  38. # )
  39. response = requests.post(
  40. "https://u475436-9425-5ad0e9a4.gda1.seetacloud.com:6443/api/v1/agent_ocr",
  41. json={
  42. "image": image_base64,
  43. }
  44. )
  45. result = response.json()
  46. print(result['data'])