forked from TDiblik/main-gate-alpr
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathutils.py
153 lines (134 loc) · 7.01 KB
/
utils.py
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
import io
import os
import shutil
import cv2
import imutils
import numpy as np
import pymssql
import concurrent
import pytesseract
from PIL import Image
from datetime import datetime
from skimage.filters import threshold_local
from ultralytics import YOLO
# Returns number of results + results as boxes
def detect_with_yolo(preloaded_model: YOLO, car_image: Image, verbose: bool) -> (int, any):
result = preloaded_model.predict(car_image, verbose=verbose)[0]
return (len(result.boxes), result.boxes)
def normalize_label(label):
return label.strip().lower()
def clean_plate_into_contours(plate_img: np.ndarray, fixed_width: int) -> np.ndarray:
# plate_img = cv2.GaussianBlur(plate_img, (5,5), 0)
plate_img = cv2.GaussianBlur(plate_img, (11,11), 0)
V = cv2.split(cv2.cvtColor(plate_img, cv2.COLOR_BGR2HSV))[2]
# T = threshold_local(V, 29, offset=15, method='gaussian')
T = threshold_local(V, 99, offset=5, method='gaussian')
thresh = (V > T).astype('uint8') * 255
thresh = cv2.bitwise_not(thresh)
plate_img = imutils.resize(plate_img, width=fixed_width)
thresh = imutils.resize(thresh, width=fixed_width)
kernel = np.ones((3, 3), np.uint8)
thresh = cv2.morphologyEx(thresh, cv2.MORPH_DILATE, kernel)
return thresh
def get_letter_rectangles_from_contours(iwl):
contours,_ = cv2.findContours(iwl,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
rectangles=[]
for cnt in contours:
x,y,w,h=cv2.boundingRect(cnt)
if(h < (iwl.shape[0] / 5) or w > (iwl.shape[1] / 5)):
continue
rectangles.append([x,y,w,h])
final_rect = []
for (x, y, w, h) in rectangles:
flag = True
for(x2, y2, w2, h2) in rectangles:
if x > x2 and y > y2 and x + w < x2 + w2 and y + h < y2 + h2:
flag = False
break
if flag:
final_rect.append([x,y,w,h])
rectangles = final_rect
rectangles.sort()
return rectangles
def gen_intermediate_file_name(filename: str, file_type: str, unique_identifier: str):
return f"./intermediate_detection_files/{filename}_{unique_identifier}.{file_type}"
# This function should be called right after getting license plate boxes
def prepare_env_for_reading_license_plates(should_save_intermediate_files: bool):
try:
if should_save_intermediate_files:
if os.path.exists("./intermediate_detection_files/"):
shutil.rmtree("./intermediate_detection_files/")
os.mkdir("./intermediate_detection_files/")
except: pass
# Read single license plate box
# Returns license plate as string
def read_license_plate(unique_identifier: str, box: any, original_image: Image, width_boost: int, additional_white_spacing_each_side: int, debug: bool, should_try_lp_crop: bool, minimum_number_of_chars_for_match: int) -> (Image, str):
# Crop image
x_min, y_min, x_max, y_max = box.xyxy.cpu().detach().numpy()[0]
original_width = x_max - x_min
original_height = y_max - y_min
boost_multiplier = width_boost / original_width
boosted_width = int(original_width * boost_multiplier)
boosted_height = int(original_height * boost_multiplier)
license_plate_cropped_img = original_image.crop( # crop to license plate
(x_min, y_min, x_max, y_max)
).convert( # black and white images make preprocessing more effective
"L"
).resize( # resizing makes recognition more effective
[boosted_width, boosted_height]
)
if should_try_lp_crop:
# crop from left and right, because license plate recognition matches with overflow
license_plate_cropped_img = license_plate_cropped_img.crop(((45, 0, boosted_width - 20, boosted_height)))
if debug:
license_plate_cropped_img.save(gen_intermediate_file_name("cropped_license_plate_full", "jpg", unique_identifier))
# Pre-process the image
license_plate_cropped_img_as_np_array = cv2.cvtColor(np.array(license_plate_cropped_img), cv2.COLOR_GRAY2BGR)
iwl_bb = clean_plate_into_contours(license_plate_cropped_img_as_np_array, width_boost)
iwl_wb = cv2.bitwise_not(iwl_bb)
if debug:
cv2.imwrite(gen_intermediate_file_name("iwl_bb", "jpg", unique_identifier), iwl_bb, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
cv2.imwrite(gen_intermediate_file_name("iwl_wb", "jpg", unique_identifier), iwl_wb, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
# Get each letter
letter_rectangles = get_letter_rectangles_from_contours(iwl_wb)
iwl_wb_pil = Image.fromarray(iwl_wb, mode="L")
# No reason to try reading, if there aren't even enought rectagles (skips reading which improves performance)
if len(letter_rectangles) < minimum_number_of_chars_for_match:
return (license_plate_cropped_img, "")
def process_letter_rectangle(args):
i, (x, y, w, h) = args
letter_box_cropped_img = iwl_wb_pil.crop((x, y, x + w, y + h))
new_letter_box_img = Image.new(
"L", (
letter_box_cropped_img.width + (additional_white_spacing_each_side * 2),
letter_box_cropped_img.height + (additional_white_spacing_each_side * 2)
), "white"
)
new_letter_box_img.paste(letter_box_cropped_img, (additional_white_spacing_each_side, additional_white_spacing_each_side))
if debug:
new_letter_box_img.save(gen_intermediate_file_name("cropped_image", "jpg", f"{unique_identifier}_{i}"))
char_from_img = pytesseract.image_to_string(new_letter_box_img, lang="eng", config="--psm 13 --dpi 96 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
char_from_img = str(char_from_img).replace("O", "0").strip()
if debug:
print(f"{i} => {char_from_img}")
return char_from_img
resulting_license_plate_string = ""
with concurrent.futures.ThreadPoolExecutor() as executor:
results = list(executor.map(process_letter_rectangle, enumerate(letter_rectangles)))
resulting_license_plate_string = "".join(results)
return (license_plate_cropped_img, resulting_license_plate_string.strip())
# https://stackoverflow.com/a/55117662/16638833
def img_to_bytes(image: Image, format="JPEG"):
bytes_io = io.BytesIO()
image.save(bytes_io, format=format)
return bytes_io.getvalue()
def save_validated_result(db_enabled: bool, car_id: str, license_plate: str, db_server: str, db_port: str, db_name: str, db_user: str, db_pass: str, save_results_enabled: bool, results_path: str, car_image_raw: Image, license_plate_image_raw: Image):
if db_enabled:
conn = pymssql.connect(server=db_server, port=db_port, database=db_name, user=db_user, password=db_pass)
cursor = conn.cursor()
cursor.execute("insert into main_gate_alpr_license_plates (id, license_plate, captured_at) values (%s, %s, %s)", (car_id, license_plate, datetime.now()))
conn.commit()
conn.close()
if save_results_enabled:
car_image_raw.save(os.path.join(results_path, f"{car_id}_car.jpg"), "JPEG")
license_plate_image_raw.save(os.path.join(results_path, f"{car_id}_lp.jpg"), "JPEG")