Spaces:
Sleeping
Sleeping
File size: 5,335 Bytes
0373fcc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
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 pdf2image import convert_from_bytes
from multiprocessing import Pool
model =YOLO("yolo_model/model_3_openvino_model")
labels = ['Achievement', 'Certifications', 'Community', 'Contact', 'Education', 'Experience', 'Interests', 'Languages', 'Name', 'Profil', 'Projects', 'image', 'resume', 'skills']
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])
y2 = max(bbox1[3], bbox2[3])
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 ):
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)
draw.rectangle([(x1, y1), (x2, y2)], outline=(255, 0, 0), width=2)
def extract_image(image,box):
x1, y1, x2, y2 = box
cropped_image = image.crop((x1, y1, x2, y2))
def process_bbox(args):
image, bbox = args
return textract_ocr(image, bbox)
def convert_bboxes_to_original(bboxes, original_size, resized_size):
"""
Convert bounding boxes from resized image size to original image size using NumPy.
:param bboxes: NumPy array of bounding boxes in format [x1, y1, x2, y2] for resized image
:param original_size: Tuple (original_width, original_height)
:param resized_size: Tuple (resized_width, resized_height)
:return: NumPy array of bounding boxes in format [x1, y1, x2, y2] for original image
"""
original_width, original_height = original_size
resized_width, resized_height = resized_size
# Calculate scaling factors
x_scale = original_width / resized_width
y_scale = original_height / resized_height
# Convert bounding boxes using broadcasting
bboxes_np = np.array(bboxes)
bboxes_np[:, 0] *= x_scale # Scale x1
bboxes_np[:, 1] *= y_scale # Scale y1
bboxes_np[:, 2] *= x_scale # Scale x2
bboxes_np[:, 3] *= y_scale # Scale y2
return bboxes_np
def correct_to_by_classifier(text):
pass
def extract_text_from_sections(image):
cv_parse = {}
original_size = image.size
original_img = image
image = image.resize((640, 640))
image = image.convert("RGB")
image_np = np.array(image)
# Perform model prediction
result = model(source=image_np, conf=0.20)
names = result[0].names # Class names
data = result[0].boxes.data.numpy()
# Extract bounding boxes and their corresponding class labels
bboxes = data[:, 0:4].tolist()
class_ids = data[:, 5].astype(int).tolist()
bboxes_filter = filter_bboxes(bboxes)
original_bboxes = convert_bboxes_to_original(bboxes_filter, original_size, (640,640))
for bbox, class_id in zip(original_bboxes, class_ids):
class_name = names[class_id]
if class_name !="image":
text = textract_ocr(original_img, bbox)
if class_name in cv_parse:
cv_parse[class_name] += "\n" + text
else:
cv_parse[class_name] = text
return cv_parse
def merge_dicts_append_strings(dict1, dict2):
# Create a new dictionary to hold the merged results
merged_dict = {}
# Add all key-value pairs from dict1 to merged_dict
for key, value in dict1.items():
merged_dict[key] = value
# Append values from dict2 to merged_dict
for key, value in dict2.items():
if key in merged_dict:
merged_dict[key] += "\n" + value
else:
merged_dict[key] = value
return merged_dict
def cv_to_json(file):
cv_parsing = {}
images = convert_from_bytes(file.read())
for image in images :
cv_parsing = merge_dicts_append_strings(cv_parsing,extract_text_from_sections(image))
return cv_parsing
|