File size: 3,148 Bytes
56ab51d |
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 |
from .extractor import CharExtractor
from ultralytics import YOLO
import numpy as np
import cv2
class CharExtraction:
def __init__(self, wPlatePath, wCharacterPath, classList, sizePlate, conf):
self.wPlatePath = wPlatePath
self.wCharacterPath = wCharacterPath
self.classList = classList
self.width, self.height = sizePlate
self.conf = conf
self.model_plat = YOLO(self.wPlatePath)
self.model_char = YOLO(self.wCharacterPath)
self.extractor = CharExtractor()
def predict(self, image, conf):
if conf == None:
self.conf = 0.3
else:
self.conf = conf
charpredict = []
platePrediction = self.model_plat(image, conf=self.conf, verbose=False) ##Melakukan prediksi dan rectange gambar
platBoxes, croppedPlateImg, confident = self._get_cropped_plate(platePrediction)
# print('lcrop : ', len(croppedPlateImg))
for croppImg in croppedPlateImg:
characterPrediction = self.model_char(croppImg, conf=self.conf, verbose=False)
charpredict.append(characterPrediction)
# print('lchar : ', len(charpredict))
characterBBOX, characterClass = self._get_bboxClass_char(charpredict, self.classList)
# print('charBox : ', len(characterBBOX))
plateNumber = self.extractor._extract_bbox(characterBBOX, characterClass)
# print('Plate Number : ', plateNumber)
# print('plateboxes : ', platBoxes)
# print('platenum : ', plateNumber)
return platBoxes, plateNumber, confident
def _get_cropped_plate(self, imgPlate): # Get cropped plate image from whole image
croppedImage = [] ##Bisa berisi banyak bbox
boxes = imgPlate[0].boxes
originalImage = imgPlate[0].orig_img
confid = boxes.conf.numpy()
h,w,c = originalImage.shape
if len(boxes) != 0:
for box in boxes.xyxy:
x1,y1,x2,y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
croppedPlateImage = originalImage[y1:y2, x1:x2, :]
croppedPlateImage = cv2.resize(croppedPlateImage, (self.width, self.height))
croppedImage.append(croppedPlateImage)
else:
croppedImage.append(originalImage)
return boxes.xyxy, croppedImage, confid
def _get_bboxClass_char(self, imgCroppedPlate, classList):
boxChar, classC = [], []
for imgCropped in imgCroppedPlate:
boxes = imgCropped[0].boxes
if len(boxes.cls) != 0:
# print(boxes.cls)
classId = [int(x) for x in boxes.cls.numpy()]
# print('classId : ', classId)
classCharacter = np.array(classList[classId])
characterBBOX = boxes.xywh
boxChar.append(characterBBOX)
classC.append(classCharacter)
# else:
# classCharacter = []
# characterBBOX = []
return boxChar, classC
|