File size: 6,217 Bytes
c1b3a0c |
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 |
import string
import easyocr
# Initialize the OCR reader
reader = easyocr.Reader(['en'], gpu=False)
# 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
|