Commit 530e3fc6 authored by sikhin.vc's avatar sikhin.vc

updated model and added logs

parent 52481e7d
...@@ -11,4 +11,4 @@ host=localhost ...@@ -11,4 +11,4 @@ host=localhost
port=8200 port=8200
[PATH] [PATH]
model_path = scripts/utility/model/yolov8l_r1_r2_lp_best.pt model_path = scripts/utility/model/r1_r2_lpr_yolov8s.pt
...@@ -55,4 +55,5 @@ uvicorn==0.19.0 ...@@ -55,4 +55,5 @@ uvicorn==0.19.0
ultralytics ultralytics
paddleocr==2.7.0.3 paddleocr==2.7.0.3
paddlepaddle==2.5.2 paddlepaddle==2.5.2
#loguru==0.7.2
...@@ -3,103 +3,113 @@ from ultralytics import YOLO ...@@ -3,103 +3,113 @@ from ultralytics import YOLO
import torch import torch
import cv2 import cv2
from scripts.config.app_configurations import * from scripts.config.app_configurations import *
from scripts.logging.application_logging import logger
device = 'cuda' if torch.cuda.is_available() else 'cpu' device = 'cuda' if torch.cuda.is_available() else 'cpu'
# print(f'Using device: {device}')
def load_ocr_model(): def load_ocr_model():
"""
Load paddle OCR model
"""
ocr = PaddleOCR(use_angle_cls=True, lang='en') ocr = PaddleOCR(use_angle_cls=True, lang='en')
logger.info("Loaded OCR model")
return ocr return ocr
def load_yolo_model(): def load_yolo_model():
"""
Load license plate detection model
"""
model = YOLO(MODEL_PATH).to(device) model = YOLO(MODEL_PATH).to(device)
logger.info("Loaded Yolo model")
return model return model
def predict(img_path, ocr, yolo_model): def predict(img_path, ocr, yolo_model):
img = cv2.imread(img_path) try:
results = yolo_model.predict(img, iou=0.7)
if results: img = cv2.imread(img_path)
r1_txt = None
r2_txt = None results = yolo_model.predict(img, iou=0.7) # predict r1, r2 and license plate using yolo model r1 - single letter region in license plate, r2 - four or five letter region in license plate
for r in results: if results:
r1_txt = None
# annotator = Annotator(img) r2_txt = None
boxes = r.boxes for r in results:
for box in boxes:
b = box.xyxy[0] # get box coordinates in (top, left, bottom, right) format
# print("bb: ",b.tolist())
x1 = int(b[0]) boxes = r.boxes
y1 = int(b[1]) for box in boxes:
x2 = int(b[2]) b = box.xyxy[0] # get box coordinates in (top, left, bottom, right) format
y2 = int(b[3]) # print("bb: ",b.tolist())
w = abs(x2 - x1) x1 = int(b[0])
h = abs(y2 - y1) y1 = int(b[1])
x2 = int(b[2])
cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 1) y2 = int(b[3])
w = abs(x2 - x1)
# lp = img[max(y1 - 30, 0):min(y1+h + 30, img.shape[0]), max(x1 - 30, 0):min(x1+w + 30, img.shape[1])] h = abs(y2 - y1)
lp = img[y1:y1 + h, x1:x1 + w]
c = box.cls lp = img[y1:y1 + h, x1:x1 + w]
# print("class: ", c) c = box.cls
# annotator.box_label(b, model.names[int(c)])
# img = annotator.result() if int(c) == 0:
# cv2.imshow('YOLO V8 Detection', cv2.resize(img, (900, 600))) #extracting text from r1
if int(c) == 0: # cv2.imwrite("LP.png", lp)
result = ocr.ocr(lp, cls=True) result = ocr.ocr(lp, cls=True)
for idx in range(len(result)): for idx in range(len(result)):
res = result[idx] res = result[idx]
text_x = 10
text_y = 30 if res:
if res:
for ind, line in enumerate(res):
for ind, line in enumerate(res):
# print("line")
# print(line) text_tuple = line[1]
if r1_txt is None:
text_tuple = line[1] r1_txt = text_tuple[0]
if r1_txt is None: logger.info(f"r1 text : {r1_txt}")
r1_txt = text_tuple[0] text_conf = text_tuple[1]
print("R1 TEXT: ", r1_txt) elif int(c) == 1:
text_conf = text_tuple[1] #extracting text from r2
elif int(c) == 1: result = ocr.ocr(lp, cls=True)
for idx in range(len(result)):
result = ocr.ocr(lp, cls=True) res = result[idx]
for idx in range(len(result)):
res = result[idx] if res:
text_x = 10
text_y = 30 for ind, line in enumerate(res):
if res:
for ind, line in enumerate(res): text_tuple = line[1]
# print("line") if r2_txt is None:
# print(line) r2_txt = text_tuple[0]
logger.info(f"r2 text : {r2_txt}")
text_tuple = line[1] text_conf = text_tuple[1]
if r2_txt is None: # cv2.putText(lp, txt, (text_x , text_y + ind*50), cv2.FONT_HERSHEY_SIMPLEX , 1, (255, 0, 0), 2, cv2.LINE_AA)
r2_txt = text_tuple[0]
print("R2 TEXT: ", r2_txt) # cv2.imshow('detected LPR', lp)
text_conf = text_tuple[1] # cv2.waitKey(0)
# cv2.putText(lp, txt, (text_x , text_y + ind*50), cv2.FONT_HERSHEY_SIMPLEX , 1, (255, 0, 0), 2, cv2.LINE_AA) #
# cv2.imshow('detected LPR', lp) if r1_txt is not None and r2_txt is not None:
# cv2.waitKey(0) LP = f"{r1_txt} {r2_txt}"
#
logger.info(f"License plate reading : {LP}")
if r1_txt is not None and r2_txt is not None: return LP
LP = f"{r1_txt} {r2_txt}"
return "Image Not Clear"
# print("LP: ", LP) except Exception as e:
return LP logger.info(f"Error while reading license plate : {e}")
return f"Error while reading license plate : {e}"
return "Image Not Clear"
# result = ocr.ocr(img_path, cls=False,det=True,rec=True)[0]
# txts = [line[1][0] for line in result]
......
...@@ -8,6 +8,7 @@ from scripts.core.handler.Text_Extraction_Handler import load_ocr_model,load_yol ...@@ -8,6 +8,7 @@ from scripts.core.handler.Text_Extraction_Handler import load_ocr_model,load_yol
from scripts.logging.application_logging import logger from scripts.logging.application_logging import logger
from scripts.config.app_constants import TEXT_EXTRACT from scripts.config.app_constants import TEXT_EXTRACT
ocr = load_ocr_model() ocr = load_ocr_model()
yolo_model = load_yolo_model() yolo_model = load_yolo_model()
router = APIRouter() router = APIRouter()
......
test.jpg

550 KB | W: | H:

test.jpg

414 KB | W: | H:

test.jpg
test.jpg
test.jpg
test.jpg
  • 2-up
  • Swipe
  • Onion skin
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment