Spaces:
Runtime error
Runtime error
File size: 1,560 Bytes
8cd053a |
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 |
from ultralytics import YOLO
import numpy as np
from PIL import Image
import cv2 as cv
class DetectionPipeline():
def __init__(self) -> None:
# Initialize and Load in Custom YOLOv8 model training weights
self.license_plate_detector = YOLO('YOLOv8_best_weights.pt')
def preprocess_image(self, uploaded_file):
"""
Takes a File in the format - '.jpg' / '.png' / '.jpeg' then converts it into a Numpy array and returns it
Args:
uploaded_file: File that needs to be converted to Numpy array
Returns:
img_array: Image in the `np.array` format
"""
img = Image.open(uploaded_file).convert('RGB')
img_array = np.array(img)
return img_array
def detectLicensePlates(self, input_array):
detections = self.license_plate_detector(input_array)[0]
license_plate_detections = []
for license_plate in detections.boxes.data.tolist():
x1, y1, x2, y2, score, class_id = license_plate
license_plate_detections.append([int(x1), int(y1), int(x2), int(y2), score])
return license_plate_detections
def detections2Image(self, preprocess_image: np.array, detections:list):
img = np.array(preprocess_image, dtype='uint8')
for license_plate_info in detections:
x1, y1, x2, y2, score = license_plate_info
cv.rectangle(img, pt1=(x1, y1), pt2=(x2, y2), color=(0, 255, 0), thickness=2)
img_detections = np.array(img)
return img_detections
|