from transformers import TrOCRProcessor, VisionEncoderDecoderModel from PIL import Image, ImageFilter import torch import re processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten") model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten") def clean_ocr_text(text): text = text.replace(",", ".").replace("s", "5").replace("o", "0").replace("O", "0") return re.sub(r"[^\d.kg]", "", text.lower()) def extract_unit_from_text(raw_text): raw = raw_text.lower() if "kg" in raw: return "kg" elif "g" in raw: return "g" return "g" # default fallback def restore_decimal(text): if re.fullmatch(r"\d{5}", text): return f"{text[:2]}.{text[2:]}" elif re.fullmatch(r"\d{4}", text): return f"{text[:2]}.{text[2:]}" return text def extract_weight(image): try: image = image.resize((image.width * 2, image.height * 2), Image.BICUBIC) image = image.filter(ImageFilter.SHARPEN) pixel_values = processor(images=image, return_tensors="pt").pixel_values generated_ids = model.generate(pixel_values) raw_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() cleaned = clean_ocr_text(raw_text) match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)?", cleaned) if match: return f"{match.group(1)} {match.group(2) or ''}", raw_text fallback = re.search(r"\d{4,5}", cleaned) if fallback: fixed = restore_decimal(fallback.group()) return f"{fixed}", raw_text return f"No valid weight found | OCR: {cleaned}", raw_text except Exception as e: return f"Error: {str(e)}", ""