diff --git a/OCR YOLO11/yolomain.py b/OCR YOLO11/yolomain.py index e501f7d..19caf18 100644 --- a/OCR YOLO11/yolomain.py +++ b/OCR YOLO11/yolomain.py @@ -1,37 +1,36 @@ from ultralytics import YOLO -from ultralytics.utils.downloads import safe_download -from ultralytics.utils.plotting import Annotator, colors from PIL import Image +from transformers import TrOCRProcessor, VisionEncoderDecoderModel import os os.chdir("C:/Users/celma/OneDrive - Hanze/School/periode 1.4/IOT/YOLO11/License Plate Recognition.v11i.yolov11") -crop_dir_name = "C:/Users/celma/OneDrive - Hanze/School/periode 1.4/IOT/YOLO11/License Plate Recognition.v11i.yolov11/test/crops" -if not os.path.exists(crop_dir_name): - os.makedirs(crop_dir_name) -model = YOLO("license_plate_detector.pt") # Loads the pre-trained YOLO11 model +model = YOLO("license_plate_detector.pt") names = model.names -results = model.predict("test/images/000c.jpg", show=True, save=True) -img = Image.open("test/images/000c.jpg") # Load the original image using PIL +results = model.predict("test/images/000i.jpg", show=False, save=False) +img = Image.open("test/images/000i.jpg") boxes = results[0].boxes.xyxy.cpu().tolist() clss = results[0].boxes.cls.cpu().tolist() -# annotator = Annotator(results[0], line_width=2, example=names) -#print(results) print(boxes) if boxes is not None: for box, cls in zip(boxes,clss): - # annotator.box_label(box, color=colors(int(cls), True), label=names(int(cls))) - crop_obj = boxes[int(box[1]):int(box[3]) + int(box[0]):int(box[2])] cropped_img = img.crop( (int(box[0]), int(box[1]), int(box[2]), int(box[3])) ) -save_path = os.path.join(crop_dir_name, "cropped_image.jpg") -cropped_img.save("cropped_image.jpg") -cropped_img.show() \ No newline at end of file +cropped_img.show() + +processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten") +model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-printed") + +image = cropped_img.convert("RGB") +pixel_values = processor(images=image, return_tensors="pt").pixel_values +generated_ids = model.generate(pixel_values) +text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] +print(text)