Main.
py
from ultralytics import YOLO
import cv2
import util
from sort.sort import *
from util import get_car, read_license_plate, write_csv
results = {}
mot_tracker = Sort()
# load models
coco_model = YOLO('yolov8n.pt')
license_plate_detector = YOLO('./models/license_plate_detector.pt')
# load video
cap = cv2.VideoCapture('./sample.mp4')
vehicles = [2, 3, 5, 7]
# read frames
frame_nmr = -1
ret = True
while ret:
frame_nmr += 1
ret, frame = cap.read()
if ret:
results[frame_nmr] = {}
# detect vehicles
detections = coco_model(frame)[0]
detections_ = []
for detection in detections.boxes.data.tolist():
x1, y1, x2, y2, score, class_id = detection
if int(class_id) in vehicles:
detections_.append([x1, y1, x2, y2, score])
# track vehicles
track_ids = mot_tracker.update(np.asarray(detections_))
# detect license plates
license_plates = license_plate_detector(frame)[0]
for license_plate in license_plates.boxes.data.tolist():
x1, y1, x2, y2, score, class_id = license_plate
# assign license plate to car
xcar1, ycar1, xcar2, ycar2, car_id = get_car(license_plate, track_ids)
if car_id != -1:
# crop license plate
license_plate_crop = frame[int(y1):int(y2), int(x1): int(x2), :]
# process license plate
license_plate_crop_gray = cv2.cvtColor(license_plate_crop,
cv2.COLOR_BGR2GRAY)
_, license_plate_crop_thresh = cv2.threshold(license_plate_crop_gray, 64, 255,
cv2.THRESH_BINARY_INV)
# read license plate number
license_plate_text, license_plate_text_score =
read_license_plate(license_plate_crop_thresh)
if license_plate_text is not None:
results[frame_nmr][car_id] = {'car': {'bbox': [xcar1, ycar1, xcar2, ycar2]},
'license_plate': {'bbox': [x1, y1, x2, y2],
'text': license_plate_text,
'bbox_score': score,
'text_score': license_plate_text_score}}
# write results
write_csv(results, './test.csv')
Util.py
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
Visualize.py
import ast
import cv2
import numpy as np
import pandas as pd
def draw_border(img, top_left, bottom_right, color=(0, 255, 0), thickness=10,
line_length_x=200, line_length_y=200):
x1, y1 = top_left
x2, y2 = bottom_right
cv2.line(img, (x1, y1), (x1, y1 + line_length_y), color, thickness) #-- top-left
cv2.line(img, (x1, y1), (x1 + line_length_x, y1), color, thickness)
cv2.line(img, (x1, y2), (x1, y2 - line_length_y), color, thickness) #-- bottom-left
cv2.line(img, (x1, y2), (x1 + line_length_x, y2), color, thickness)
cv2.line(img, (x2, y1), (x2 - line_length_x, y1), color, thickness) #-- top-right
cv2.line(img, (x2, y1), (x2, y1 + line_length_y), color, thickness)
cv2.line(img, (x2, y2), (x2, y2 - line_length_y), color, thickness) #-- bottom-right
cv2.line(img, (x2, y2), (x2 - line_length_x, y2), color, thickness)
return img
results = pd.read_csv('./test_interpolated.csv')
# load video
video_path = 'sample.mp4'
cap = cv2.VideoCapture(video_path)
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Specify the codec
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
out = cv2.VideoWriter('./out.mp4', fourcc, fps, (width, height))
license_plate = {}
for car_id in np.unique(results['car_id']):
max_ = np.amax(results[results['car_id'] == car_id]['license_number_score'])
license_plate[car_id] = {'license_crop': None,
'license_plate_number': results[(results['car_id'] == car_id) &
(results['license_number_score'] == max_)]
['license_number'].iloc[0]}
cap.set(cv2.CAP_PROP_POS_FRAMES, results[(results['car_id'] == car_id) &
(results['license_number_score'] == max_)]
['frame_nmr'].iloc[0])
ret, frame = cap.read()
x1, y1, x2, y2 = ast.literal_eval(results[(results['car_id'] == car_id) &
(results['license_number_score'] == max_)]
['license_plate_bbox'].iloc[0].replace('[ ', '[').replace(' ', ' ').replace(' ', ' ').replace(' ', ','))
license_crop = frame[int(y1):int(y2), int(x1):int(x2), :]
license_crop = cv2.resize(license_crop, (int((x2 - x1) * 400 / (y2 - y1)), 400))
license_plate[car_id]['license_crop'] = license_crop
frame_nmr = -1
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
# read frames
ret = True
while ret:
ret, frame = cap.read()
frame_nmr += 1
if ret:
df_ = results[results['frame_nmr'] == frame_nmr]
for row_indx in range(len(df_)):
# draw car
car_x1, car_y1, car_x2, car_y2 = ast.literal_eval(df_.iloc[row_indx]
['car_bbox'].replace('[ ', '[').replace(' ', ' ').replace(' ', ' ').replace(' ', ','))
draw_border(frame, (int(car_x1), int(car_y1)), (int(car_x2), int(car_y2)), (0, 255,
0), 25,
line_length_x=200, line_length_y=200)
# draw license plate
x1, y1, x2, y2 = ast.literal_eval(df_.iloc[row_indx]
['license_plate_bbox'].replace('[ ', '[').replace(' ', ' ').replace(' ', ' ').replace(' ', ','))
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 12)
# crop license plate
license_crop = license_plate[df_.iloc[row_indx]['car_id']]['license_crop']
H, W, _ = license_crop.shape
try:
frame[int(car_y1) - H - 100:int(car_y1) - 100,
int((car_x2 + car_x1 - W) / 2):int((car_x2 + car_x1 + W) / 2), :] =
license_crop
frame[int(car_y1) - H - 400:int(car_y1) - H - 100,
int((car_x2 + car_x1 - W) / 2):int((car_x2 + car_x1 + W) / 2), :] = (255,
255, 255)
(text_width, text_height), _ = cv2.getTextSize(
license_plate[df_.iloc[row_indx]['car_id']]['license_plate_number'],
cv2.FONT_HERSHEY_SIMPLEX,
4.3,
17)
cv2.putText(frame,
license_plate[df_.iloc[row_indx]['car_id']]['license_plate_number'],
(int((car_x2 + car_x1 - text_width) / 2), int(car_y1 - H - 250 +
(text_height / 2))),
cv2.FONT_HERSHEY_SIMPLEX,
4.3,
(0, 0, 0),
17)
except:
pass
out.write(frame)
frame = cv2.resize(frame, (1280, 720))
# cv2.imshow('frame', frame)
# cv2.waitKey(0)
out.release()
cap.release()
Add_missing_data.py
import csv
import numpy as np
from scipy.interpolate import interp1d
def interpolate_bounding_boxes(data):
# Extract necessary data columns from input data
frame_numbers = np.array([int(row['frame_nmr']) for row in data])
car_ids = np.array([int(float(row['car_id'])) for row in data])
car_bboxes = np.array([list(map(float, row['car_bbox'][1:-1].split())) for row in data])
license_plate_bboxes = np.array([list(map(float, row['license_plate_bbox'][1:-
1].split())) for row in data])
interpolated_data = []
unique_car_ids = np.unique(car_ids)
for car_id in unique_car_ids:
frame_numbers_ = [p['frame_nmr'] for p in data if int(float(p['car_id'])) ==
int(float(car_id))]
print(frame_numbers_, car_id)
# Filter data for a specific car ID
car_mask = car_ids == car_id
car_frame_numbers = frame_numbers[car_mask]
car_bboxes_interpolated = []
license_plate_bboxes_interpolated = []
first_frame_number = car_frame_numbers[0]
last_frame_number = car_frame_numbers[-1]
for i in range(len(car_bboxes[car_mask])):
frame_number = car_frame_numbers[i]
car_bbox = car_bboxes[car_mask][i]
license_plate_bbox = license_plate_bboxes[car_mask][i]
if i > 0:
prev_frame_number = car_frame_numbers[i-1]
prev_car_bbox = car_bboxes_interpolated[-1]
prev_license_plate_bbox = license_plate_bboxes_interpolated[-1]
if frame_number - prev_frame_number > 1:
# Interpolate missing frames' bounding boxes
frames_gap = frame_number - prev_frame_number
x = np.array([prev_frame_number, frame_number])
x_new = np.linspace(prev_frame_number, frame_number, num=frames_gap,
endpoint=False)
interp_func = interp1d(x, np.vstack((prev_car_bbox, car_bbox)), axis=0,
kind='linear')
interpolated_car_bboxes = interp_func(x_new)
interp_func = interp1d(x, np.vstack((prev_license_plate_bbox,
license_plate_bbox)), axis=0, kind='linear')
interpolated_license_plate_bboxes = interp_func(x_new)
car_bboxes_interpolated.extend(interpolated_car_bboxes[1:])
license_plate_bboxes_interpolated.extend(interpolated_license_plate_bboxes[1:])
car_bboxes_interpolated.append(car_bbox)
license_plate_bboxes_interpolated.append(license_plate_bbox)
for i in range(len(car_bboxes_interpolated)):
frame_number = first_frame_number + i
row = {}
row['frame_nmr'] = str(frame_number)
row['car_id'] = str(car_id)
row['car_bbox'] = ' '.join(map(str, car_bboxes_interpolated[i]))
row['license_plate_bbox'] = ' '.join(map(str,
license_plate_bboxes_interpolated[i]))
if str(frame_number) not in frame_numbers_:
# Imputed row, set the following fields to '0'
row['license_plate_bbox_score'] = '0'
row['license_number'] = '0'
row['license_number_score'] = '0'
else:
# Original row, retrieve values from the input data if available
original_row = [p for p in data if int(p['frame_nmr']) == frame_number and
int(float(p['car_id'])) == int(float(car_id))][0]
row['license_plate_bbox_score'] = original_row['license_plate_bbox_score'] if
'license_plate_bbox_score' in original_row else '0'
row['license_number'] = original_row['license_number'] if 'license_number' in
original_row else '0'
row['license_number_score'] = original_row['license_number_score'] if
'license_number_score' in original_row else '0'
interpolated_data.append(row)
return interpolated_data
# Load the CSV file
with open('test.csv', 'r') as file:
reader = csv.DictReader(file)
data = list(reader)
# Interpolate missing data
interpolated_data = interpolate_bounding_boxes(data)
# Write updated data to a new CSV file
header = ['frame_nmr', 'car_id', 'car_bbox', 'license_plate_bbox',
'license_plate_bbox_score', 'license_number', 'license_number_score']
with open('test_interpolated.csv', 'w', newline='') as file:
writer = csv.DictWriter(file, fieldnames=header)
writer.writeheader()
writer.writerows(interpolated_data)