File size: 5,663 Bytes
0a2675c |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
from model1 import reader, np, YOLO, car_detection, lp_detection
import torch
from PIL import Image
import cv2
from torchvision import transforms
char_dect = YOLO("models/yolov8n_lpchar_det.pt")
char_rec = torch.load("models/charrec.pt", map_location="cpu")
# function to detect cars in the given image
def detect_cars(inputs):
cars = []
# running the cars detection model with 50% confidence threshold
car_results = car_detection.predict(source=inputs, classes=[2], conf=0.5, verbose=False)
# iterating through each output (num of outputs will be same as num of inputs)
for car_result in car_results:
# finding the bounding boxes of the cars detected
boxes = car_result.boxes.xyxy.tolist()
# iterating through each car detected
for box in boxes:
# cropping car image from the input image
car = car_result.orig_img[int(box[1]):int(box[3]), int(box[0]):int(box[2])]
cars.append(car)
return cars
# function to detect licence plates in the given car images
def detect_lp(inputs):
lps = []
# running the license plate detection model with 50% confidence threshold
lp_results = lp_detection.predict(source=inputs, conf=0.5, verbose=False)
# iterating through each output (num of outputs will be same as num of inputs)
for lp_result in lp_results:
# finding the bounding boxes of the license plate detected
lp_boxes = lp_result.boxes.xyxy.tolist()
# iterating through each license plate detected
for lp_box in lp_boxes:
# cropping license plate image from the car image
lp = lp_result.orig_img[int(lp_box[1]):int(lp_box[3]), int(lp_box[0]):int(lp_box[2])]
lps.append(lp)
# breaking as we only want to detect one licence plate per car
break
# if no licence plate is detected then we are adding a black image
if len(lp_boxes) == 0:
lps.append(np.zeros((100,100,3), np.uint8))
return lps
# function to detect licence plates character in the given LP images
def chars_lp_det(inputs):
vis_lp = []
chars = []
# running the license plate detection model with 50% confidence threshold
chars_results = char_dect.predict(source=inputs, conf=0.5, verbose=False)
# iterating through each output (num of outputs will be same as num of inputs)
for chars_result in chars_results:
# finding the bounding boxes of the license plate detected
chars_boxes = chars_result.boxes.xyxy.tolist()
# iterating through each license plate detected
vis = chars_result.orig_img.copy()
c_list =[]
for chars_box in chars_boxes:
# cropping license plate image from the car image
cv2.rectangle(vis, (int(chars_box[0]),int(chars_box[1])), (int(chars_box[2]), int(chars_box[3])), (0,255,0), 1)
chrs = chars_result.orig_img[int(chars_box[1]):int(chars_box[3]), int(chars_box[0]):int(chars_box[2])]
c_list.append(chrs)
chars.append(c_list)
vis_lp.append(vis)
# if no licence plate is detected then we are adding a black image
if len(vis_lp) == 0:
vis_lp.append(np.zeros((100,100,3), np.uint8))
return vis_lp, chars
# function to detect licence plate number in the given licence plate images
def detect_lp_text(inputs):
plate_number = []
# iterating through each licence plate
for input in inputs:
# finding the number/text in licence plate
result = reader.readtext(input)
# if no text is found in the licence plate, then adding a default text not found
if len(result) == 0:
plate_number.append("not found")
else:
# adding the licence plate number to a list
plate_number.append(result[0][1])
return plate_number
def rec_lp_char(inputs):
m = ['0','1','2','3','4','5','6','7','8','9','A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
lptexts = []
for input in inputs:
imgs = [transform(Image.fromarray(input[i])) for i in range(len(input))]
if len(imgs) <= 1:
lptexts.append("not found")
continue
imgs = torch.stack(imgs)
output = char_rec(imgs)
preds = torch.argmax(output, dim=1).tolist()
lptext = ""
for pred in preds:
lptext += m[int(pred)]
lptexts.append(lptext)
return lptexts
def run(inputs):
# for future, to handle multiple inputs
# currently using just one input
inputs = inputs[0]
# detecting cars, this function returns all detected car images
cars = detect_cars(inputs)
# if no car is detected black images are returned
if len(cars) == 0:
return [np.zeros((100,100,3), np.uint8)], [np.zeros((100,100,3), np.uint8)], "not found"
# detecting licence plates from the car images
# returns licence plate images, if it cant find a license plate a black image is returned
lps = detect_lp(cars)
vis_lp, chars_lp = chars_lp_det(lps)
lptexts = rec_lp_char(chars_lp)
# detecting licence plate number from licence plate images
# returns text from the licence plate images, if none is detected "not found" text is returned
# lp_text = detect_lp_text(lps)
return cars, vis_lp, lptexts
|