|
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() |
|
|
|
|
|
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_ids = mot_tracker.update(np.asarray(detections_)) |
|
|
|
license_plates = plate_model(img)[0] |
|
for license_plate in license_plates.boxes.data.tolist(): |
|
x1, y1, x2, y2, score, class_id = license_plate |
|
|
|
xcar1, ycar1, xcar2, ycar2, car_id = get_car(license_plate, track_ids) |
|
if car_id != -1: |
|
|
|
license_plate_crop = img[int(y1):int(y2), int(x1): int(x2), :] |
|
|
|
text = plate_rec(license_plate_crop) |
|
return text, license_plate_crop |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
