File size: 1,352 Bytes
ffbdc6a |
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 |
from ultralytics import YOLO
import cv2
from sort.sort import *
from util import get_car, read_license_plate, write_csv
from PIL import Image
from micro_ocr import plate_rec
def plate_reader(plate_model,yolo_model,img_path):
mot_tracker = Sort()
# load models
vehicles = [2, 3, 5, 7]
img = cv2.imread(img_path)
detections = yolo_model(img)[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 = plate_model(img)[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 = img[int(y1):int(y2), int(x1): int(x2), :]
text = plate_rec(license_plate_crop)
return text, license_plate_crop
|