Update yolo_text_extraction.py

yolo_text_extraction.py

@@ -1,32 +1,24 @@
-from ultralytics import YOLO 
-from PIL import Image,ImageDraw
-import numpy as np 
+from ultralytics import YOLO
+from PIL import Image, ImageDraw
+import numpy as np
 from PIL import ImageFilter
-
 from dotenv import load_dotenv
-
-import numpy as np
-from ocr_functions import paddle_ocr,textract_ocr,tesseract_ocr
+from ocr_functions import paddle_ocr, textract_ocr, tesseract_ocr
 from pdf2image import convert_from_path
 
-
-model =YOLO("yolo_model/best.pt")
-
+model = YOLO("yolo_model/best.pt")
 
 def check_intersection(bbox1, bbox2):
-    # Check for intersection between two bounding boxes
     x1, y1, x2, y2 = bbox1
     x3, y3, x4, y4 = bbox2
     return not (x3 > x2 or x4 < x1 or y3 > y2 or y4 < y1)
 
 def check_inclusion(bbox1, bbox2):
-    # Check if one bounding box is completely inside another
     x1, y1, x2, y2 = bbox1
     x3, y3, x4, y4 = bbox2
     return x1 >= x3 and y1 >= y3 and x2 <= x4 and y2 <= y4
 
 def union_bbox(bbox1, bbox2):
-    # Calculate the union of two bounding boxes
     x1 = min(bbox1[0], bbox2[0])
     y1 = min(bbox1[1], bbox2[1])
     x2 = max(bbox1[2], bbox2[2])
@@ -34,43 +26,32 @@ def union_bbox(bbox1, bbox2):
     return [x1, y1, x2, y2]
 
 def filter_bboxes(bboxes):
-    # Iterate through each pair of bounding boxes and filter out those that intersect or are completely contained within another
     filtered_bboxes = []
     for bbox1 in bboxes:
         is_valid = True
         for bbox2 in filtered_bboxes:
             if check_intersection(bbox1, bbox2):
-                # If the two bounding boxes intersect, compute their union
                 bbox1 = union_bbox(bbox1, bbox2)
-                # Mark the current bbox as invalid to be removed
                 is_valid = False
                 break
             elif check_inclusion(bbox1, bbox2):
-                # If bbox1 is completely contained within bbox2, mark bbox1 as invalid to be removed
                 is_valid = False
                 break
         if is_valid:
             filtered_bboxes.append(bbox1)
     return filtered_bboxes
 
-
-
-
-def draw_bboxes(image, bboxes ):
+def draw_bboxes(image, bboxes):
     draw = ImageDraw.Draw(image)
     for bbox in bboxes:
         x1, y1, x2, y2 = bbox
-        
-        x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
         draw.rectangle([(x1, y1), (x2, y2)], outline=(255, 0, 0), width=2)
 
-
-
-def extract_image(image,box):
+def extract_image(image, box):
     x1, y1, x2, y2 = box
     cropped_image = image.crop((x1, y1, x2, y2))
 
-
 def text_image(image):
     image = image.convert("RGB")
     image = image.filter(ImageFilter.MedianFilter(3))
@@ -86,13 +67,11 @@ def text_image(image):
     draw_bboxes(image, bboxes_filter)
     image.save("output.png")
     texts = [textract_ocr(image, bbox) for bbox in bboxes_filter]
-    return "\n------section-------\n"+"\n------section-------\n".join(texts)
-
-    
+    return "\n------section-------\n" + "\n------section-------\n".join(texts)
 
 def pdf_to_text(pdf_file):
     text = ""
     images = convert_from_path(pdf_file)
-    for image in images : 
+    for image in images:
         text = text + text_image(image) + "\n"
-    return text 
+    return text