AIFilterContent / services /image_ocr.py
Mridul2003's picture
Update services/image_ocr.py
44f72ba verified
raw
history blame
702 Bytes
from PIL import Image
from io import BytesIO
import requests
from transformers import pipeline
class ImageOCRService:
def __init__(self):
self.pipe = pipeline("image-text-to-text", model="ds4sd/SmolDocling-256M-preview")
def extract_text(self, image_url: str) -> str:
response = requests.get(image_url)
image = Image.open(BytesIO(response.content)).convert("RGB")
result = self.pipe([{
"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": "extract text from image"}
]
}])
return result[0]['generated_text'][1]['content'] if result else ""