python -m venv venv
번호판 인식 pre-trained ai 깃허브에서 zip 다운로드
해당 디렉토리로 이동해서 pip install
pip3 install requirements.txt
예시 코드
pip install easyocr
pip install yolov5
pip3 install torch torchvision torchaudio --index-url <https://download.pytorch.org/whl/cu118>
# 같은 디렉토리에 image.jpg 파일 있어야함
from glob import glob
import torch
import easyocr
import numpy as np
import cv2
import sys
from PIL import ImageFont, ImageDraw, Image
def load_model():
car_m = torch.hub.load('./yolov5', 'yolov5s', './yolov5/models/yolo.py', force_reload=True, skip_validation=True, source='local')
lp_m = torch.hub.load('./yolov5', 'custom', 'lp_det.pt', source='local')
reader = easyocr.Reader(['en'], detect_network='craft', recog_network='best_acc', user_network_directory='lp_models/user_network', model_storage_directory='lp_models/models')
car_m.classes = [2,3, 5, 7]
return car_m, lp_m, reader
def main():
car_m, lp_m, reader = load_model()
file = "./image.jpg"
im, text = detect(car_m, lp_m, reader, file)
file = open("result.txt", "a", encoding= 'utf8')
for i in text:
file.write(i + "\\n")
file.close()
def detect(car_m, lp_m, reader, path):
fontpath = "SpoqaHanSansNeo-Light.ttf"
font = ImageFont.truetype(fontpath, 200)
im = Image.open(path)
im.save(path)
to_draw = np.array(im)
results = car_m(im)
locs = results.xyxy[0]
result_text = []
if len(locs) == 0:
result = lp_m(im)
if len(result) == 0:
result_text.append('Car was not detected')
else:
for rslt in result.xyxy[0]:
x2,y2,x3,y3 = [item1.cpu().detach().numpy().astype(np.int32) for item1 in rslt[:4]]
try:
extra_boxes = 0
im = cv2.cvtColor(cv2.resize(to_draw[y2 - extra_boxes:y3 + extra_boxes,x2 - extra_boxes:x3 + extra_boxes], (224,128)), cv2.COLOR_BGR2GRAY)
text = reader.recognize(im)[0][1]
result_text.append(text)
except Exception as e:
return cv2.resize(to_draw, (1280,1280)), ""
img_pil = Image.fromarray(to_draw)
draw = ImageDraw.Draw(img_pil)
draw.text( (x2-100,y2-300), text, font=font, fill=(255,0,255))
to_draw = np.array(img_pil)
cv2.rectangle(to_draw, (x2.item(),y2.item()),(x3.item(),y3.item()),(255,0,255),10)
return cv2.resize(to_draw, (1280,1280)), result_text
return cv2.resize(to_draw, (1280,1280)), result_text
if __name__ == '__main__':
main()