Implemented easy ocr

api/routers/reply_generator.py

@@ -16,7 +16,7 @@ class UserRequest(BaseModel):
     tones: Optional[List[str]] = None
 
 @router.post("/reply-generator")
-async def generate_reply(request: UserRequest):
+def generate_reply(request: UserRequest):
     """
     Endpoint to extract conversation from a base64-encoded image
     """

requirements.txt

@@ -7,4 +7,7 @@ requests
 fastapi
 uvicorn 
 python-dotenv
-streamlit
+streamlit
+easyocr
+opencv-python
+pillow

src/conversation_extractor/extractor.py

@@ -1,11 +1,15 @@
 
 from google import genai
 from google.genai import types
-from src.utils.models_loader import ocr_llm
+from src.utils.models_loader import ocr_llm , reply_llm , reader
 from .prompts import prompt
-
+import easyocr
+from PIL import Image
+import io
+import numpy as np
 from dotenv import load_dotenv
 import os
+from langchain_core.messages import HumanMessage , SystemMessage
 
 load_dotenv()
 os.environ['GOOGLE_API_KEY']=os.getenv('GOOGLE_API_KEY')
@@ -15,18 +19,42 @@ class ConversationExtractor:
         self.client = genai.Client()
         self.model_name = ocr_llm
         self.prompt = prompt
+    
+    def complete_ocr(self, image_bytes:bytes , lang_list=['en']):
+        img = Image.open(io.BytesIO(image_bytes))
+        w_img = img.width
+
+        # Convert PIL image to RGB and read it directly using EasyOCR
+        results = reader.readtext(np.array(img))  # Use numpy array instead of file path
+
+        conversation = []
+        for bbox, text, conf in results:
+            # bbox: [[x1,y1],[x2,y2],[x3,y3],[x4,y4]]
+            x_coords = [p[0] for p in bbox]
+            y_coords = [p[1] for p in bbox]
+            x, y, w, h = int(min(x_coords)), int(min(y_coords)), int(max(x_coords)-min(x_coords)), int(max(y_coords)-min(y_coords))
 
-    def extract_conversation(self, image_bytes: bytes) -> str:
-        """
-        Extract conversation text from an image.
-        :param image_bytes: Binary content of the image
-        :return: Extracted conversation text
-        """
-        response = self.client.models.generate_content(
-            model=self.model_name,
-            contents=[
-                types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg"),
-                self.prompt
-            ]
+            # Determine left or right speaker
+            speaker = "A" if (x + w/2) < w_img/2 else "B"
+
+            conversation.append({
+                "speaker": speaker,
+                "text": text.strip(),
+                "box": [x, y, w, h]
+            })
+
+        # Sort top to bottom
+        raw_ocr = sorted(conversation, key=lambda x: x["box"][1])
+        raw_ocr_text = "Detected Conversation:\n" + "\n".join(
+        [f"Line {i}: {turn['text']}" for i, turn in enumerate(raw_ocr, start=1)]
         )
-        return response.text
+        return raw_ocr_text
+    
+    def extract_conversation(self, image_bytes:bytes):
+        raw_ocr_text=self.complete_ocr(image_bytes)
+        messages = [SystemMessage(content = prompt),
+        HumanMessage(content = raw_ocr_text)]
+        response =reply_llm.invoke(messages)
+        print('The cleaned ocr:', response.content)
+        return response.content
+

src/conversation_extractor/prompts.py

@@ -1,8 +1,27 @@
+prompt = '''
+You are a text cleaner and conversation extractor.
+You will be given raw OCR text from a screenshot of a chat. The text may contain:
 
-prompt = (
-    "Extract only the conversation text between two persons from this image. "
-    "Ignore timestamps, emojis, or extra elements. "
-    "Preserve dialogue order. "
-    "If names exist, use them; otherwise label speakers Person A and Person B. "
-    "Output strictly as chat dialogue lines. Return only dialogue."
-)
+- Timestamps (e.g., "05:32 PM", "05.33 PM")
+- System messages of mobile phone (e.g., "New Contact", "New Message")
+- Battery percentage, signal info, or other UI elements
+- Numbers or non-message text
+
+Your task is to extract **only the actual conversation messages exchanged between the two people**, in the **order they appear**.
+
+**Rules:**
+1. Remove all system messages, timestamps, numbers, and noise.
+2. Keep the text messages only.
+3. Keep the order of messages intact.
+"
+
+
+**Output Format:**
+A clean list of messages between the two people, one message per line. Format like this:
+
+Line 1: Hello!
+Line 2: Hi, how are you?
+Line 3: I'm good, thanks! And you?
+...
+
+'''

src/utils/models_loader.py

@@ -1,8 +1,14 @@
 from langchain_groq import ChatGroq
 from dotenv import load_dotenv
 import os
+import easyocr
 
 load_dotenv()
 os.environ['GROQ_API_KEY']=os.getenv('GROQ_API_KEY')
 reply_llm = ChatGroq(model='llama-3.1-8b-instant')
-ocr_llm = "gemini-2.5-flash"
+ocr_llm = "gemini-2.5-flash"
+
+
+
+lang_list = ['en']
+reader = easyocr.Reader(lang_list, gpu=True)