File size: 4,663 Bytes
3f118e2 |
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 |
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() |