Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| from ultralytics.engine.results import Results | |
| import cv2 | |
| import os | |
| import numpy as np | |
| class Recognize_ID: | |
| def __init__(self,model_path='../../recognization_id.pt'): | |
| current_dir = os.path.dirname(os.path.abspath(__file__)) | |
| model_path = os.path.join(current_dir , model_path ) | |
| self.model = YOLO(model=model_path ) | |
| def give_me_id_number(self,image:str): | |
| """ | |
| image_dir : input image directory | |
| model : yolo model | |
| """ | |
| if isinstance(image, str): # If the input is a file path | |
| current_dir = os.path.dirname(os.path.abspath(__file__)) | |
| image_path = os.path.join(current_dir , image ) | |
| img = cv2.imread(image_path) | |
| elif isinstance(image, np.ndarray): # If the input is a NumPy array | |
| img = image | |
| res = self.model(img) | |
| boxes = res[0].boxes.xywh[::,0].tolist() | |
| classes = res[0].boxes.cls.tolist() | |
| boxes_labels =[(int(key) , int(value)) for key , value in zip(boxes, classes)] | |
| boxes_labels.sort() | |
| national_id = "".join([str(i[1]) for i in boxes_labels]) | |
| return national_id | |