Spaces:
Sleeping
Sleeping
| 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 | |