0% found this document useful (0 votes)
12 views

Code

Uploaded by

lp3859826
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as DOCX, PDF, TXT or read online on Scribd
0% found this document useful (0 votes)
12 views

Code

Uploaded by

lp3859826
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as DOCX, PDF, TXT or read online on Scribd
You are on page 1/ 18

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)

You might also like