Update utility/utils.py

utility/utils.py

@@ -96,10 +96,11 @@ def process_image(image_path, scale=2):
     return final_image
 
 # Function for OCR with PaddleOCR, returning both text and bounding boxes
-def ocr_with_paddle(img):
+def ocr_with_paddle(img_path):
     final_text = ''
-    
-    logging.info(f"PADDLEOCR_HOME: {os.environ['PADDLEOCR_HOME']}")
+    boxes = []
+
+    # Initialize PaddleOCR
     ocr = PaddleOCR(
         lang='en', 
         use_angle_cls=True,
@@ -108,21 +109,33 @@ def ocr_with_paddle(img):
         cls_model_dir=os.path.join(os.environ['PADDLEOCR_HOME'], 'whl/cls/ch_ppocr_mobile_v2.0_cls_infer')
     )
 
-    result = ocr.ocr(img)
-    boxes = []
+    # Perform OCR on the image
+    result = ocr.ocr(img_path)
+
+    # Load image with OpenCV
+    img = cv2.imread(img_path)
+
+    # Iterate through OCR results
     for line in result[0]:
         box, text, _ = line
-        boxes.append(box)  # Append the bounding box
         final_text += ' ' + text
+        boxes.append(box)  # Save the bounding box coordinates
+
+        # Convert points to integer and draw the bounding box
+        points = [(int(point[0]), int(point[1])) for point in box]
+        cv2.polylines(img, [np.array(points)], isClosed=True, color=(0, 255, 0), thickness=2)
 
-    return final_text, boxes
+    # Save the image with drawn boxes in memory (as a variable)
+    img_with_boxes = img  # This image can be used later or saved to disk if needed
+
+    return final_text, img_with_boxes
 
 # Function to draw bounding boxes around text
-def draw_boxes(image, boxes):
-    draw = ImageDraw.Draw(image)
-    for box in boxes:
-        draw.polygon(box, outline="red", width=3)
-    return image
+#def draw_boxes(image, boxes):
+#    draw = ImageDraw.Draw(image)
+#    for box in boxes:
+#        draw.polygon(box, outline="red", width=3)
+#    return image
 
 # Extract text and create a result image with bounding boxes
 def extract_text_from_images(image_paths):
@@ -134,11 +147,11 @@ def extract_text_from_images(image_paths):
             enhanced_image = process_image(image_path, scale=2)
 
             # Perform OCR on the enhanced image and get boxes
-            result, boxes = ocr_with_paddle(enhanced_image)
+            result, img_with_boxes = ocr_with_paddle(enhanced_image)
 
             # Draw bounding boxes on the processed image
             img_result = Image.fromarray(enhanced_image)
-            img_with_boxes = draw_boxes(img_result, boxes)
+            #img_with_boxes = draw_boxes(img_result, boxes)
 
             # Save the image with boxes
             result_image_path = os.path.join(RESULT_FOLDER, f'result_{os.path.basename(image_path)}')