Update yolomain.py

Werkende OCR zit er nu in.
Nutteloze lines zijn weggehaald.
show image kan nog weg eventueel
This commit is contained in:
zoloko13
2025-06-12 10:51:40 +02:00
parent 78854fa84f
commit ea0669df61

View File

@@ -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()
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)