from transformers import TrOCRProcessor, VisionEncoderDecoderModel from PIL import Image, ImageFilter import torch import re # Load model processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten") model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten") def clean_ocr_text(text): print("[RAW OCR]", text) text = text.replace(",", ".").replace("s", "5").replace("o", "0").replace("O", "0") text = re.sub(r"[^\d.kg]", "", text.lower()) # Keep digits, dots, and kg print("[CLEANED OCR]", text) return text 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) # Try direct match: e.g., 52.25 kg or 75.0 g match = re.search(r"(\d{1,3}\.\d{1,3})\s*(kg|g)", cleaned) if match: return f"{match.group(1)} {match.group(2)}" # Try fallback: extract digits and manually guess decimal fallback_match = re.search(r"(\d{4,5})", cleaned) if fallback_match: fallback_value = restore_decimal(fallback_match.group(1)) # Check for presence of unit hints in raw_text unit = "kg" if "kg" in raw_text.lower() else "g" return f"{fallback_value} {unit}" return f"No valid weight found | OCR: {cleaned}" except Exception as e: return f"Error: {str(e)}"