Update app.py

app.py

@@ -1,94 +1,90 @@
-import streamlit as st
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
-from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import JSONResponse
-import uvicorn
-import numpy as np
 import cv2
-import io
-
-# Create a FastAPI app instance
-app = FastAPI()
+import numpy as np
+from io import BytesIO
+import uvicorn
 
 # Load the model and processor
-@st.cache_resource
-def load_model():
-    processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-    model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
-    return processor, model
-
-processor, model = load_model()
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
 
+app = FastAPI()
 
-# Function to preprocess image and detect lines (used for multiline OCR)
+# Helper function to preprocess the image and detect lines
 def detect_lines(image, min_height=20, min_width=100):
+    """
+    Detects lines of text in the given image.
+    """
+    # Convert the PIL image to a NumPy array
     image_np = np.array(image)
+
+    # Convert to grayscale
     gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
+
+    # Apply binary thresholding
     _, binary = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+
+    # Dilate to merge nearby text
     kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
     dilated = cv2.dilate(binary, kernel, iterations=1)
-    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    bounding_boxes = sorted([cv2.boundingRect(c) for c in contours], key=lambda b: b[1])
-    line_images = [image_np[y:y+h, x:x+w] for x, y, w, h in bounding_boxes if h >= min_height and w >= min_width]
-    return line_images
-
-
-# FastAPI endpoint to handle image processing
-@app.post("/process_image")
-async def process_image(image: UploadFile = File(...)):
-    try:
-        # Read the uploaded image
-        image_data = await image.read()
-        image = Image.open(io.BytesIO(image_data)).convert("RGB")
 
-        # Detect lines and process each line
-        line_images = detect_lines(image, min_height=30, min_width=600)
-        extracted_text = ""
-        for line_img in line_images:
-            line_pil = Image.fromarray(line_img)
-            pixel_values = processor(images=line_pil, return_tensors="pt").pixel_values
-            generated_ids = model.generate(pixel_values)
-            generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-            extracted_text += generated_text + "\n"
+    # Find contours
+    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
-        # Return extracted text as JSON
-        return JSONResponse(content={"extracted_text": extracted_text.strip()})
+    # Sort contours top-to-bottom
+    bounding_boxes = [cv2.boundingRect(c) for c in contours]
+    bounding_boxes = sorted(bounding_boxes, key=lambda b: b[1])  # Sort by y-coordinate
 
-    except Exception as e:
-        return JSONResponse(content={"error": str(e)}, status_code=500)
+    # Filter out small contours and merge nearby ones
+    filtered_boxes = []
+    for x, y, w, h in bounding_boxes:
+        if h >= min_height and w >= min_width:  # Filter small boxes
+            filtered_boxes.append((x, y, w, h))
 
+    # Extract individual lines as images
+    line_images = []
+    for (x, y, w, h) in filtered_boxes:
+        line = image_np[y:y + h, x:x + w]
+        line_images.append(line)
 
-# Streamlit UI for testing (optional)
-st.title("OCR API Service with Multiline Support")
+    return line_images
 
-uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
-if uploaded_file is not None:
+@app.post("/process_image")
+async def process_image(file: UploadFile = File(...)):
+    """
+    API endpoint to process the uploaded image and extract multiline text.
+    """
     try:
-        image = Image.open(uploaded_file).convert("RGB")
-        st.image(image, caption="Uploaded Image", use_column_width=True)
+        # Read the uploaded image
+        contents = await file.read()
+        image = Image.open(BytesIO(contents)).convert("RGB")
 
         # Detect lines in the image
-        st.write("Detecting lines...")
         line_images = detect_lines(image, min_height=30, min_width=100)
-        st.write(f"Detected {len(line_images)} lines in the image.")
 
         # Perform OCR on each detected line
         extracted_text = ""
         for idx, line_img in enumerate(line_images):
+            # Convert the line image to PIL format
             line_pil = Image.fromarray(line_img)
+
+            # Prepare the image for OCR
             pixel_values = processor(images=line_pil, return_tensors="pt").pixel_values
+
+            # Generate text from the line image
             generated_ids = model.generate(pixel_values)
-            generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-            extracted_text += f"{generated_text}\n"
+            line_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+            # Append the extracted text
+            extracted_text += f"{line_text}\n"
 
-        # Display extracted text
-        st.subheader("Extracted Text:")
-        st.text_area("Output Text", extracted_text.strip(), height=300)
+        # Return the extracted text as a JSON response
+        return {"extracted_text": extracted_text}
 
     except Exception as e:
-        st.error(f"An error occurred: {e}")
+        raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")
 
-# Run the FastAPI app
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)