Update app.py

extract both the serial number (Sr. No) and meter reading from the uploaded meter images.

app.py

@@ -2,12 +2,21 @@ import gradio as gr
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
 import torch
+import re
 
 # Load the pre-trained model and processor
 processor = TrOCRProcessor.from_pretrained('microsoft/trocr-small-printed')
 model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-small-printed')
 model.eval()
 
+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]*([\w\d\-]+)', text, re.IGNORECASE)
+    meter_reading = re.search(r'(?:Reading|Meter Reading|Reading:)[:\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 extract_meter_reading(image):
     try:
         # Convert and preprocess the image
@@ -19,7 +28,9 @@ def extract_meter_reading(image):
             generated_ids = model.generate(pixel_values)
         generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        return f"Extracted Meter Reading: {generated_text.strip()}"
+        # Extract details from the generated text
+        sr_no, reading = extract_details(generated_text)
+        return f"Serial Number: {sr_no}\nMeter Reading: {reading}"
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -28,8 +39,8 @@ iface = gr.Interface(
     fn=extract_meter_reading,
     inputs=gr.Image(type="numpy", label="Upload or Capture Meter Image"),
     outputs="text",
-    title="Meter Reading Extractor",
-    description="Upload a meter image to extract the reading using a lightweight OCR model."
+    title="Meter Reading and Serial Number Extractor",
+    description="Upload a meter image to extract the serial number and meter reading using a lightweight OCR model."
 )
 
 # Launch the Gradio app