Spaces:
Sleeping
Sleeping
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 | |