Code
Code
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_))
if car_id != -1:
# write results
write_csv(results, './test.csv')
Util.py
import string
import easyocr
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'))
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
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)
if license_complies_format(text):
return format_license(text), score
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]
import cv2
import numpy as np
import pandas as pd
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)
return img
results = pd.read_csv('./test_interpolated.csv')
# load video
video_path = 'sample.mp4'
cap = cv2.VideoCapture(video_path)
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()
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)
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)
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:
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]
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]))
interpolated_data.append(row)
return interpolated_data