Update app.py

app.py

@@ -1,8 +1,7 @@
 import streamlit as st
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
-import cv2
-import numpy as np
+import json
 
 # Load the model and processor
 @st.cache_resource
@@ -13,68 +12,53 @@ def load_model():
 
 processor, model = load_model()
 
-# Helper function to preprocess the image and detect lines
-def detect_lines(image, min_height=20, min_width=100):
-    # Convert the PIL image to a NumPy array
-    image_np = np.array(image)
+# Check if the request is an API call
+if st.runtime.scriptrunner.script_run_context.is_running_with_auth:
+    import io
+    from fastapi import FastAPI, File, UploadFile
+    from fastapi.responses import JSONResponse
 
-    # Convert to grayscale
-    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
+    app = FastAPI()
 
-    # Apply binary thresholding
-    _, binary = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+    @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")
 
-    # Dilate to merge nearby text
-    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
-    dilated = cv2.dilate(binary, kernel, iterations=1)
-
-    # Find contours
-    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-    # 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
-
-    # 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)
+            # Perform OCR
+            pixel_values = processor(images=image, return_tensors="pt").pixel_values
+            generated_ids = model.generate(pixel_values)
+            generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-    return line_images
+            # Return extracted text as JSON
+            return JSONResponse(content={"extracted_text": generated_text})
 
-# Streamlit app
-st.title("OCR API Service with Multiline Support")
+        except Exception as e:
+            return JSONResponse(content={"error": str(e)}, status_code=500)
+else:
+    # Streamlit UI for manual testing
+    st.title("OCR API Service")
 
-# Handle image upload
-uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+    uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
 
-if uploaded_file is not None:
-    try:
-        # Load and process the uploaded image
-        image = Image.open(uploaded_file).convert("RGB")
-        line_images = detect_lines(image, min_height=30, min_width=100)
+    if uploaded_file is not None:
+        try:
+            # Load and display the uploaded image
+            image = Image.open(uploaded_file).convert("RGB")
+            st.image(image, caption="Uploaded Image", use_column_width=True)
 
-        # Perform OCR on each detected line
-        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
+            # Perform OCR
+            pixel_values = processor(images=image, 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 += f"{generated_text}\n"
 
-        # Simulate API-like JSON response
-        json_response = {"extracted_text": extracted_text}
+            # Display extracted text
+            st.subheader("Extracted Text:")
+            st.text(generated_text)
 
-        # Return JSON response
-        st.write(json_response)  # This is the response to your CodeIgniter client
-    except Exception as e:
-        # Return an error response
-        st.write({"error": str(e)})
+        except Exception as e:
+            st.error(f"An error occurred: {e}")
+    else:
+        st.info("Please upload an image to start the OCR process.")