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