Update app.py

app.py

@@ -10,57 +10,72 @@ import numpy as np
 # Initialize EasyOCR Reader
 reader = easyocr.Reader(['en'])
 
-# Load the pre-trained model and processor
+# Load the pre-trained model and processor for fallback
 processor = TrOCRProcessor.from_pretrained('microsoft/trocr-small-printed')
 model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-small-printed')
 model.eval()
 
 def enhance_image(image):
-    # Enhance the image for better OCR results
     image = ImageEnhance.Contrast(image).enhance(2.0)
+    image = ImageEnhance.Sharpness(image).enhance(2.0)
     image = ImageOps.grayscale(image)
     return image
 
 
-def extract_details(text):
-    # Use regex to extract serial number and meter reading
-    sr_no = re.search(r'(?:Sr\.No|SR\.NO|Serial Number|SR NO)[:\s]*([\d]+)', text, re.IGNORECASE)
-    meter_reading = re.search(r'(?:Reading|Meter Reading|kWh|KW H)[:\s]*([\d]+)', text, re.IGNORECASE)
-    sr_no_text = sr_no.group(1) if sr_no else "Not Found"
-    reading_text = meter_reading.group(1) if meter_reading else "Not Found"
-    return sr_no_text, reading_text
+def ocr_with_easyocr(pil_img):
+    # Convert PIL to bytes for EasyOCR
+    buf = BytesIO()
+    pil_img.save(buf, format='PNG')
+    return reader.readtext(buf.getvalue(), detail=0)
 
 
-def extract_meter_reading(image):
-    try:
-        # Convert NumPy array to PIL Image
-        image = Image.fromarray(image)
-        enhanced_image = enhance_image(image)
-
-        # Convert PIL image to bytes for EasyOCR
-        img_byte_arr = BytesIO()
-        enhanced_image.save(img_byte_arr, format='PNG')
-        img_byte_arr = img_byte_arr.getvalue()
+def ocr_with_trocr(pil_img):
+    pixel_values = processor(images=pil_img.convert("RGB"), return_tensors="pt").pixel_values
+    with torch.no_grad():
+        ids = model.generate(pixel_values)
+    return processor.batch_decode(ids, skip_special_tokens=True)[0]
 
-        # Use EasyOCR to get better results
-        result = reader.readtext(img_byte_arr)
-        text = " ".join([item[1] for item in result])
 
-        # Extract details from the OCR text
-        sr_no, reading = extract_details(text)
-        return f"Serial Number: {sr_no}\nMeter Reading: {reading}"
+def extract_meter_reading(image):
+    try:
+        pil = Image.fromarray(image)
+        w, h = pil.size
+        # Define regions: reading top 40%, serial bottom 40%
+        top_region = pil.crop((0, 0, w, int(h * 0.4)))
+        bottom_region = pil.crop((0, int(h * 0.5), w, h))
+        # Enhance regions
+        top_enh = enhance_image(top_region)
+        bot_enh = enhance_image(bottom_region)
+        # OCR regions
+        top_texts = ocr_with_easyocr(top_enh) + [ocr_with_trocr(top_enh)]
+        bot_texts = ocr_with_easyocr(bot_enh) + [ocr_with_trocr(bot_enh)]
+        top_combined = " ".join(top_texts)
+        bot_combined = " ".join(bot_texts)
+        # Extract reading: look for digits near kwh
+        reading_match = re.search(r"(\d+)\s*(?=kwh|kw h|k w h)", top_combined, re.IGNORECASE)
+        if not reading_match:
+            # fallback: any 4-7 digit number
+            fallback = re.findall(r"\b\d{4,7}\b", top_combined)
+            reading = fallback[0] if fallback else "Not Found"
+        else:
+            reading = reading_match.group(1)
+        # Extract serial: longest digit sequence in bottom
+        serials = re.findall(r"\d{6,12}", bot_combined)
+        serial = max(serials, key=len) if serials else "Not Found"
+        return f"Serial Number: {serial}\nMeter Reading: {reading}"
     except Exception as e:
         return f"Error: {str(e)}"
 
+# Gradio app
+def main():
+    iface = gr.Interface(
+        fn=extract_meter_reading,
+        inputs=gr.Image(type="numpy", label="Upload or Capture Meter Image"),
+        outputs="text",
+        title="Meter Reading and Serial Number Extractor",
+        description="Upload a meter image; extracts serial number and meter reading using region-based OCR."
+    )
+    iface.launch()
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=extract_meter_reading,
-    inputs=gr.Image(type="numpy", label="Upload or Capture Meter Image"),
-    outputs="text",
-    title="Meter Reading and Serial Number Extractor",
-    description="Upload a meter image to extract the serial number and meter reading using OCR techniques."
-)
-
-# Launch the Gradio app
-iface.launch()
+if __name__ == '__main__':
+    main()