detect and read license plate

.gitignore

@@ -0,0 +1,4 @@
+.idea
+*.iml
+enhanced.*
+vision/__pycache__/*

detect.py

@@ -0,0 +1,28 @@
+import getopt
+import sys
+
+from vision.enhancer import enhance
+from vision.vision import detect_license_plate, read_license_plate_tesseract, \
+    read_license_plate_ml
+
+if __name__ == "__main__":
+    opts, args = getopt.getopt(sys.argv[1:], "hi:", ["input="])
+    input_url = ""
+    for opt, arg in opts:
+        if opt == '-h':
+            print('detect.py -i input_image_url')
+            sys.exit()
+        elif opt in ("-i", "--input"):
+            input_url = arg
+
+    print("detecting license plate from image: " + input_url)
+
+    crop = detect_license_plate(input_url)
+    enhanced_path = "./enhanced.png"
+
+    with enhance(crop) as enhanced:
+        enhanced.save(enhanced_path)
+
+    print("license plate trocr: " + read_license_plate_ml(enhanced_path))
+    print("license plate tesseract: " + read_license_plate_tesseract(enhanced_path))
+

requirements.txt

@@ -0,0 +1,8 @@
+pillow
+transformers
+yolov5
+pytesseract
+opencv-python
+numpy
+torch
+sympy

vision/enhancer.py

@@ -0,0 +1,73 @@
+import cv2
+import numpy as np
+from PIL import ImageOps, ImageEnhance, Image
+
+
+def enhance(crop: Image.Image) -> Image.Image:
+    """
+    prepare crop for ocr - grayscale, contrast, sharpness, resize
+    :param crop: image to prepare
+    :rtype: Image.Image
+    :return: processed image
+    """
+    image = ImageOps.grayscale(crop)
+    image = ImageEnhance.Contrast(image).enhance(5)
+    image = ImageEnhance.Sharpness(image).enhance(5)
+    image = ImageOps.scale(image, 3.5)
+
+    image = _remove_dark_frame(image)
+    image = _remove_shades(image)
+    return image
+
+
+def _remove_dark_frame(img: Image.Image) -> Image.Image:
+    # Convert PIL Image to OpenCV format
+    open_cv_image = np.array(img.convert('RGB'))
+    open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)
+
+    # Convert to grayscale
+    gray = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2GRAY)
+
+    # Blur the image to reduce noise
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+
+    # Apply a binary threshold after blurring
+    _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+
+    # Find contours from the binary image
+    cnts = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    cnts = cnts[0] if len(cnts) == 2 else cnts[1]
+
+    # Assume the largest contour is the frame
+    c = max(cnts, key=cv2.contourArea)
+
+    # Find the bounding box coordinates from the contour
+    x, y, w, h = cv2.boundingRect(c)
+
+    # Crop the original image using the bounding box coordinates
+    cropped = open_cv_image[y:y+h, x:x+w]
+
+    # Convert back to PIL format
+    cropped_pil = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+    return cropped_pil
+
+def _remove_shades(image: Image.Image, threshold: int = 64) -> Image.Image:
+    """
+    Remove shades between white and black from an image using PIL.
+
+    Parameters:
+    - image_path: The path to the image.
+    - threshold: The cutoff for determining whether a pixel should be white or black.
+
+    Returns:
+    - A new PIL Image object with shades removed.
+    """
+    # Convert the image to grayscale
+    gray_img = image.convert("L")
+
+    # Apply the threshold
+    # All pixels value > threshold will be set to 255 (white)
+    # All pixels value <= threshold will be set to 0 (black)
+    binary_img = gray_img.point(lambda x: 255 if x > threshold else 0, '1')
+
+    return binary_img

vision/vision.py

@@ -0,0 +1,44 @@
+import yolov5
+from PIL import Image
+from pytesseract import pytesseract
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+
+
+def detect_license_plate(image_path: str, save_crops=False, save_dir="/tmp/crops") -> Image.Image:
+    """
+    Detect license plate from image_path
+    :param image_path: image_path path
+    :param save_crops: save crops to filesystem
+    :param save_dir: directory to save crops
+    :rtype: :py:class:`~PIL.Image.Image`
+    :returns: An :py:class:`~PIL.Image.Image` object.
+    """
+    # load model
+    model = yolov5.load('keremberke/yolov5n-license-plate')
+
+    # set model parameters
+    model.conf = 0.25  # NMS confidence threshold
+    model.iou = 0.45  # NMS IoU threshold
+    model.agnostic = False  # NMS class-agnostic
+    model.multi_label = False  # NMS multiple labels per box
+    model.max_det = 1  # maximum number of detections per image
+
+    # perform inference
+    results = model(image_path, size=640)
+    crops = results.crop(save=save_crops, save_dir=save_dir)
+
+    return Image.fromarray(crops[0]['im'][..., ::-1])
+
+
+def read_license_plate_tesseract(image_path) -> str:
+    return pytesseract.image_to_string(image_path,
+                                       config='--psm 7 -c tessedit_char_whitelist=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ')
+
+
+def read_license_plate_ml(image_path: str) -> str:
+    processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-printed')
+    model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-printed')
+    pixel_values = processor(images=Image.open(image_path).convert("RGB"), return_tensors="pt").pixel_values
+
+    generated_ids = model.generate(pixel_values)
+    return processor.batch_decode(generated_ids, skip_special_tokens=True)[0]