mertbozkurt's picture
Update util.py
5868efa
import string
# Mapping dictionaries for character conversion
dict_char_to_int = {'O': '0',
'I': '1',
'J': '3',
'A': '4',
'G': '6',
'S': '5'}
dict_int_to_char = {'0': 'O',
'1': 'I',
'3': 'J',
'4': 'A',
'6': 'G',
'5': 'S'}
def write_csv(results, output_path):
"""
Write the results to a CSV file.
Args:
results (dict): Dictionary containing the results.
output_path (str): Path to the output CSV file.
"""
with open(output_path, 'w') as f:
f.write('{},{},{},{},{},{},{}\n'.format('frame_nmr', 'car_id', 'car_bbox',
'license_plate_bbox', 'license_plate_bbox_score', 'license_number',
'license_number_score'))
for frame_nmr in results.keys():
for car_id in results[frame_nmr].keys():
print(results[frame_nmr][car_id])
if 'car' in results[frame_nmr][car_id].keys() and \
'license_plate' in results[frame_nmr][car_id].keys() and \
'text' in results[frame_nmr][car_id]['license_plate'].keys():
f.write('{},{},{},{},{},{},{}\n'.format(frame_nmr,
car_id,
'[{} {} {} {}]'.format(
results[frame_nmr][car_id]['car']['bbox'][0],
results[frame_nmr][car_id]['car']['bbox'][1],
results[frame_nmr][car_id]['car']['bbox'][2],
results[frame_nmr][car_id]['car']['bbox'][3]),
'[{} {} {} {}]'.format(
results[frame_nmr][car_id]['license_plate']['bbox'][0],
results[frame_nmr][car_id]['license_plate']['bbox'][1],
results[frame_nmr][car_id]['license_plate']['bbox'][2],
results[frame_nmr][car_id]['license_plate']['bbox'][3]),
results[frame_nmr][car_id]['license_plate']['bbox_score'],
results[frame_nmr][car_id]['license_plate']['text'],
results[frame_nmr][car_id]['license_plate']['text_score'])
)
f.close()
def license_complies_format(text):
"""
Check if the license plate text complies with the required format.
Args:
text (str): License plate text.
Returns:
bool: True if the license plate complies with the format, False otherwise.
"""
if len(text) != 7:
return False
if (text[0] in string.ascii_uppercase or text[0] in dict_int_to_char.keys()) and \
(text[1] in string.ascii_uppercase or text[1] in dict_int_to_char.keys()) and \
(text[2] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'] or text[2] in dict_char_to_int.keys()) and \
(text[3] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'] or text[3] in dict_char_to_int.keys()) and \
(text[4] in string.ascii_uppercase or text[4] in dict_int_to_char.keys()) and \
(text[5] in string.ascii_uppercase or text[5] in dict_int_to_char.keys()) and \
(text[6] in string.ascii_uppercase or text[6] in dict_int_to_char.keys()):
return True
else:
return False
def format_license(text):
"""
Format the license plate text by converting characters using the mapping dictionaries.
Args:
text (str): License plate text.
Returns:
str: Formatted license plate text.
"""
license_plate_ = ''
mapping = {0: dict_int_to_char, 1: dict_int_to_char, 4: dict_int_to_char, 5: dict_int_to_char, 6: dict_int_to_char,
2: dict_char_to_int, 3: dict_char_to_int}
for j in [0, 1, 2, 3, 4, 5, 6]:
if text[j] in mapping[j].keys():
license_plate_ += mapping[j][text[j]]
else:
license_plate_ += text[j]
return license_plate_
def read_license_plate(license_plate_crop):
"""
Read the license plate text from the given cropped image.
Args:
license_plate_crop (PIL.Image.Image): Cropped image containing the license plate.
Returns:
tuple: Tuple containing the formatted license plate text and its confidence score.
"""
detections = reader.readtext(license_plate_crop)
for detection in detections:
bbox, text, score = detection
text = text.upper().replace(' ', '')
if license_complies_format(text):
return format_license(text), score
return None, None
def get_car(license_plate, vehicle_track_ids):
"""
Retrieve the vehicle coordinates and ID based on the license plate coordinates.
Args:
license_plate (tuple): Tuple containing the coordinates of the license plate (x1, y1, x2, y2, score, class_id).
vehicle_track_ids (list): List of vehicle track IDs and their corresponding coordinates.
Returns:
tuple: Tuple containing the vehicle coordinates (x1, y1, x2, y2) and ID.
"""
x1, y1, x2, y2, score, class_id = license_plate
foundIt = False
for j in range(len(vehicle_track_ids)):
xcar1, ycar1, xcar2, ycar2, car_id = vehicle_track_ids[j]
if x1 > xcar1 and y1 > ycar1 and x2 < xcar2 and y2 < ycar2:
car_indx = j
foundIt = True
break
if foundIt:
return vehicle_track_ids[car_indx]
return -1, -1, -1, -1, -1