File size: 2,219 Bytes
792db51 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import os
import requests
import base64
import json
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
# Load secrets from environment variables (configure these in your Hugging Face Space)
API_KEY = os.environ.get("API_KEY")
OCR_API_URL = os.environ.get("OCR_API_URL")
OCR_TASK = os.environ.get("OCR_TASK")
OCR_FUNCTION_NAME = os.environ.get("OCR_FUNCTION_NAME")
def call_ocr(image_path: str):
# Read and encode the image to base64
with open(image_path, "rb") as image_file:
base64_string = base64.b64encode(image_file.read()).decode('utf-8')
payload = {
"image": base64_string,
"task": OCR_TASK,
"function_name": OCR_FUNCTION_NAME
}
headers = {
"Content-Type": "application/json",
"Accept": "application/json",
"Authorization": f"Basic {API_KEY}"
}
response = requests.post(OCR_API_URL, json=payload, headers=headers)
response.raise_for_status()
response_json = response.json()
# Extract quad_boxes and labels
data = response_json.get('data', {}).get(OCR_TASK, {})
quad_boxes = data.get('quad_boxes', [])
labels = data.get('labels', [])
# Clean up labels
cleaned_labels = [label.replace("</s>", "").strip() for label in labels]
# Open the original image and draw boxes
image = Image.open(image_path)
draw = ImageDraw.Draw(image)
# Use a true-type font if available
try:
font = ImageFont.truetype("arial.ttf", 16)
except IOError:
font = None
for quad, label in zip(quad_boxes, cleaned_labels):
polygon = [(quad[i], quad[i+1]) for i in range(0, len(quad), 2)]
draw.polygon(polygon, outline="red", width=2)
draw.text((quad[0], quad[1]), label, fill="blue", font=font)
output_labels = "\n".join(cleaned_labels)
return image, output_labels
iface = gr.Interface(
fn=call_ocr,
inputs=gr.Image(type="filepath", label="Upload Image"),
outputs=[gr.Image(label="Annotated Image"), gr.Textbox(label="Detected Labels")],
title="Optical Character Recognition (OCR)"
)
if __name__ == "__main__":
iface.launch(debug=True)
|