YOLOv8でナンバープレートを検出してAI-OCRで読み取ってみました。 〜ファインチューニングに使用したデータは、撮影した写真ではなく、Pythonで生成した画像（30,000枚・192,000アノテーション）です〜

"""
ナンバープレート画像を生成するクラス

number_plate = NumberPlate()

img = number_plate.generate(category)

category: 分類
0:自家用普通車
1:事業用普通車
2:自家用軽自動車
3:事業用軽自動車
"""
import random
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont


class NumberPlate:
    def __init__(self):
        self.__fontFace_0 = "./TrmFontJB.ttf"
        self.__fontFace_1 = "./hiiragi_marugo_ProN_W4.ttc"

        black = (0, 0, 0)
        green = (0, 60, 0)
        white = (240, 240, 240)
        yellow = (0, 240, 240)
        gray = (128, 128, 128)

        self.__height = 220
        self.__width = 440
        self.__margin = 4

        self.__back_color_list = [white, green, yellow, black]
        self.__text_color_list = [green, white, black, yellow]
        self.__gray_color = gray

    def __drawText(
        self, img, text, position, fontFace, fontScale, text_color, stroke_width
    ):
        (r, g, b) = text_color
        color = (b, g, r)
        # cv2 -> PIL
        imgPIL = Image.fromarray(cv2.cvtColor(img.copy(), cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(imgPIL)
        fontPIL = ImageFont.truetype(font=fontFace, size=fontScale)
        draw.text(
            xy=position, text=text, fill=color, font=fontPIL, stroke_width=stroke_width
        )
        # PIL -> cv2
        return cv2.cvtColor(np.array(imgPIL, dtype=np.uint8), cv2.COLOR_RGB2BGR)

    def __generate_place(self):
        list = [
            "札幌",
            "釧路",
            "山形",
            "水戸",
            "群馬",
            "長野",
            "大阪",
            "京都",
            "姫路",
            "多摩",
            "埼玉",
            "室蘭",
            "千葉",
            "釧路",
            "川崎",
        ]

        index = random.randint(0, len(list) - 1)

        return list[index]

    def __generate_classification(self):
        list = [
            "500",
            "300",
            "551",
            "512",
            "330",
            "331",
            "50",
            "30",
            "55",
            "51",
            "33",
            "31",
        ]
        index = random.randint(0, len(list) - 1)
        return list[index]

    def __generate_hiragana(self):
        list = ["あ", "き", "い", "を", "う", "く", "か"]
        index = random.randint(0, len(list) - 1)
        return list[index]

    def __generate_number(self):
        number = random.randint(1, 9999)
        isHead = True

        y = int(number / 1000)
        if isHead and y == 0:
            number_text = "・"
        else:
            isHead = False
            number_text = str(y)
        number -= y * 1000

        y = int(number / 100)
        if isHead and y == 0:
            number_text += "・"
        else:
            isHead = False
            number_text += str(y)
        number -= y * 100

        if isHead:
            number_text += " "
        else:
            if random.randint(0, 3) != 0:
                number_text += "-"
            else:
                number_text += " "

        y = int(number / 10)
        if isHead and y == 0:
            number_text += "・"
        else:
            number_text += str(y)

        number -= y * 10
        number_text += str(number)

        return number_text

    def generate(self, category):
        img = np.full((self.__height, self.__width, 3), 0, dtype=np.uint8)
        cv2.rectangle(
            img,
            (0, 0),
            (self.__width, self.__height),
            self.__back_color_list[category],
            thickness=-1,
        )
        cv2.rectangle(
            img,
            (0, 0),
            (self.__width - 2, self.__height - 2),
            self.__gray_color,
            thickness=1,
        )
        cv2.rectangle(
            img,
            (self.__margin, self.__margin),
            (self.__width - self.__margin * 2, self.__height - self.__margin * 2),
            self.__gray_color,
            thickness=1,
        )

        position = (100, 10)
        fontScale = 55
        stroke_width = 1
        img = self.__drawText(
            img,
            self.__generate_place(),
            position,
            self.__fontFace_1,
            fontScale,
            self.__text_color_list[category],
            stroke_width,
        )

        position = (230, 10)
        fontScale = 65
        stroke_width = 0
        img = self.__drawText(
            img,
            self.__generate_classification(),
            position,
            self.__fontFace_0,
            fontScale,
            self.__text_color_list[category],
            stroke_width,
        )
        position = (20, 80)
        fontScale = 150
        stroke_width = 0
        img = self.__drawText(
            img,
            self.__generate_hiragana(),
            position,
            self.__fontFace_0,
            fontScale,
            self.__text_color_list[category],
            stroke_width,
        )
        position = (80, 80)
        fontScale = 130
        stroke_width = 0
        img = self.__drawText(
            img,
            self.__generate_number(),
            position,
            self.__fontFace_0,
            fontScale,
            self.__text_color_list[category],
            stroke_width,
        )

        radius = 6
        position = (80, 30)
        img = cv2.circle(
            img, position, radius, self.__gray_color, thickness=-1, lineType=cv2.LINE_AA
        )
        position = (360, 30)
        img = cv2.circle(
            img, position, radius, self.__gray_color, thickness=-1, lineType=cv2.LINE_AA
        )
        return img

"""
ナンバープレートを生成しアフィン変換して背景と合成することでGround Truth形式のデータセットを作成する
"""
import json
import glob
import random
import os
import shutil
import math
import numpy as np
import cv2
from PIL import Image
from number_plate import NumberPlate

MAX = 30000  # 生成する画像数

CLASS_NAME = [
    "PRIVATE_STANDARD",
    "BUS_STANDAR",
    "PRIVATE_LIGHT_VEHICLE",
    "BUS_LIGHT_VEHICLE",
]

COLORS = [(0, 0, 175), (175, 0, 0), (0, 175, 0), (175, 0, 175)]

BACKGROUND_IMAGE_PATH = "./dataset/background_images"
TARGET_IMAGE_PATH = "./dataset/output_png"
OUTPUT_PATH = "./dataset/output_ground_truth"

S3Bucket = "s3://ground_truth_dataset"
manifestFile = "output.manifest"


BASE_WIDTH = 440  # ナンバープレートの基本サイズは、背景画像とのバランスより、横幅を440を基準とする
BACK_WIDTH = 1280  # 背景画像ファイルのサイズを合わせる必要がある
BACK_HEIGHT = 960  # 背景画像ファイルのサイズを合わせる必要がある


# 背景画像取得クラス
class Background:
    def __init__(self, backPath):
        self.__backPath = backPath

    def get(self):
        imagePath = random.choice(glob.glob(self.__backPath + "/*.png"))
        return cv2.imread(imagePath, cv2.IMREAD_UNCHANGED)


# 検出対象取得クラス (base_widthで指定された横幅を基準にリサイズされる)
class Target:
    def __init__(self, target_path, base_width, class_name):
        self.__target_path = target_path
        self.__base_width = base_width
        self.__class_name = class_name
        self.__number_plate = NumberPlate()

    def get(self, class_id):
        # 対象画像
        class_name = self.__class_name[class_id]
        target_image = number_plate = self.__number_plate.generate(class_id)

        # 透過PNGへの変換
        # Point 1: 白色部分に対応するマスク画像を生成
        mask = np.all(target_image[:, :, :] == [255, 255, 255], axis=-1)
        # Point 2: 元画像をBGR形式からBGRA形式に変換
        target_image = cv2.cvtColor(target_image, cv2.COLOR_BGR2BGRA)
        # Point3: マスク画像をもとに、白色部分を透明化
        target_image[mask, 3] = 0

        # 基準（横）サイズに基づきリサイズ
        h, w, _ = target_image.shape
        aspect = h / w
        target_image = cv2.resize(
            target_image, (int(self.__base_width * aspect), self.__base_width)
        )

        return target_image


# 変換クラス
class Transformer:
    def __init__(self, width, height):
        self.__width = width
        self.__height = height
        self.__min_scale = 0.1
        self.__max_scale = 1

    # distonation UP:0,DOWN;1,LEFT:2,RIGHT:3
    # ratio: UP/DOWN 0 〜 0.1
    # ratio: LEFT/RIGHT 0 〜 0.3
    def __affine(self, target_image, ratio, distnatoin):
        # アフィン変換（上下、左右から見た傾き）
        h, w, _ = target_image.shape
        if distnatoin == 0:
            distortion_w = int(ratio * w)
            x1 = 0
            x2 = 0 - distortion_w
            x3 = w
            x4 = w - distortion_w
            y1 = h
            y2 = 0
            y3 = 0
            y4 = h
            # 変換後のイメージのサイズ
            ww = w + distortion_w
            hh = h
        elif distnatoin == 1:
            distortion_w = int(ratio * w)
            x1 = 0 - distortion_w
            x2 = 0
            x3 = w - distortion_w
            x4 = w
            y1 = h
            y2 = 0
            y3 = 0
            y4 = h
            # 変換後のイメージのサイズ
            ww = w + distortion_w
            hh = h
        elif distnatoin == 2:
            distortion_h = int(ratio * h)
            x1 = 0
            x2 = 0
            x3 = w
            x4 = w
            y1 = h
            y2 = 0 - int(distortion_h * 0.6)
            y3 = 0
            y4 = h - distortion_h
            # 変換後のイメージのサイズ
            ww = w
            hh = h + int(distortion_h * 1.3)
        elif distnatoin == 3:
            distortion_h = int(ratio * h)
            x1 = 0
            x2 = 0
            x3 = w
            x4 = w
            y1 = h - int(distortion_h * 0.6)
            y2 = 0
            y3 = 0 - distortion_h
            y4 = h
            # 変換後のイメージのサイズ
            ww = w
            hh = h + int(distortion_h * 1.3)

        pts2 = [(x2, y2), (x1, y1), (x4, y4), (x3, y3)]
        w2 = max(pts2, key=lambda x: x[0])[0]
        h2 = max(pts2, key=lambda x: x[1])[1]
        h, w, _ = target_image.shape
        pts1 = np.float32([(0, 0), (0, h), (w, h), (w, 0)])
        pts2 = np.float32(pts2)

        M = cv2.getPerspectiveTransform(pts2, pts1)
        target_image = cv2.warpPerspective(
            target_image, M, (w2 + 100, h2 + 100), borderValue=(255, 255, 255)
        )
        return (target_image, ww, hh)

    def warp(self, target_image):
        # サイズ変更
        target_image = self.__resize(target_image)

        # アフィン変換
        # distonation UP:0,DOWN;1,LEFT:2,RIGHT:3
        # ratio: UP/DOWN max:0.1
        # ratio: LEFT/RIGHT max:0.3
        ratio = random.uniform(0, 0.1)
        distonation = random.randint(0, 3)
        if distonation == 2 or distonation == 3:
            ratio = random.uniform(0, 0.3)
        (target_image, w, h) = self.__affine(target_image, ratio, distonation)

        # 配置位置決定
        left = random.randint(0, self.__width - w)
        top = random.randint(0, self.__height - h)
        rect = ((left, top), (left + w, top + h))

        # 背景面との合成
        new_image = self.__synthesize(target_image, left, top)
        return (new_image, rect)

    def __resize(self, img):
        scale = random.uniform(self.__min_scale, self.__max_scale)
        w, h, _ = img.shape
        return cv2.resize(img, (int(w * scale), int(h * scale)))

    def __rote(self, target_image, angle):
        h, w, _ = target_image.shape
        rate = h / w
        scale = 1
        if rate < 0.9 or 1.1 < rate:
            scale = 0.9
        elif rate < 0.8 or 1.2 < rate:
            scale = 0.6
        center = (int(w / 2), int(h / 2))
        trans = cv2.getRotationMatrix2D(center, angle, scale)
        return cv2.warpAffine(target_image, trans, (w, h))

    def __synthesize(self, target_image, left, top):
        background_image = np.zeros((self.__height, self.__width, 4), np.uint8)
        back_pil = Image.fromarray(background_image)
        front_pil = Image.fromarray(target_image)
        back_pil.paste(front_pil, (left, top), front_pil)
        return np.array(back_pil)


class Effecter:
    # Gauss
    def gauss(self, img, level):
        return cv2.blur(img, (level * 2 + 1, level * 2 + 1))

    # Noise
    def noise(self, img):
        img = img.astype("float64")
        img[:, :, 0] = self.__single_channel_noise(img[:, :, 0])
        img[:, :, 1] = self.__single_channel_noise(img[:, :, 1])
        img[:, :, 2] = self.__single_channel_noise(img[:, :, 2])
        return img.astype("uint8")

    def __single_channel_noise(self, single):
        diff = 255 - single.max()
        noise = np.random.normal(0, random.randint(1, 100), single.shape)
        noise = (noise - noise.min()) / (noise.max() - noise.min())
        noise = diff * noise
        noise = noise.astype(np.uint8)
        dst = single + noise
        return dst


# バウンディングボックス描画
def box(frame, rect, class_id):
    ((x1, y1), (x2, y2)) = rect
    label = "{}".format(CLASS_NAME[class_id])
    img = cv2.rectangle(frame, (x1, y1), (x2, y2), COLORS[class_id], 2)
    img = cv2.rectangle(img, (x1, y1), (x1 + 150, y1 - 20), COLORS[class_id], -1)
    cv2.putText(
        img,
        label,
        (x1 + 2, y1 - 2),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.5,
        (255, 255, 255),
        1,
        cv2.LINE_AA,
    )
    return img


# 背景と商品の合成
def marge_image(background_image, front_image):
    back_pil = Image.fromarray(background_image)
    front_pil = Image.fromarray(front_image)
    back_pil.paste(front_pil, (0, 0), front_pil)
    return np.array(back_pil)


# Manifest生成クラス
class Manifest:
    def __init__(self, class_name):
        self.__lines = ""
        self.__class_map = {}
        for i in range(len(class_name)):
            self.__class_map[str(i)] = class_name[i]

    def appned(self, fileName, data, height, width):
        date = "0000-00-00T00:00:00.000000"
        line = {
            "source-ref": "{}/{}".format(S3Bucket, fileName),
            "boxlabel": {
                "image_size": [{"width": width, "height": height, "depth": 3}],
                "annotations": [],
            },
            "boxlabel-metadata": {
                "job-name": "xxxxxxx",
                "class-map": self.__class_map,
                "human-annotated": "yes",
                "objects": {"confidence": 1},
                "creation-date": date,
                "type": "groundtruth/object-detection",
            },
        }
        for i in range(data.max()):
            (_, rect, class_id) = data.get(i)
            ((x1, y1), (x2, y2)) = rect
            line["boxlabel"]["annotations"].append(
                {
                    "class_id": class_id,
                    "width": x2 - x1,
                    "top": y1,
                    "height": y2 - y1,
                    "left": x1,
                }
            )
        self.__lines += json.dumps(line) + "\n"

    def get(self):
        return self.__lines


# 1画像分のデータを保持するクラス
class Data:
    def __init__(self, rate):
        self.__rects = []
        self.__images = []
        self.__class_ids = []
        self.__rate = rate

    def get_class_ids(self):
        return self.__class_ids

    def max(self):
        return len(self.__rects)

    def get(self, i):
        return (self.__images[i], self.__rects[i], self.__class_ids[i])

    # 追加（重複率が指定値以上の場合は失敗する）
    def append(self, target_image, rect, class_id):
        conflict = False
        for i in range(len(self.__rects)):
            iou = self.__multiplicity(self.__rects[i], rect)
            if iou > self.__rate:
                conflict = True
                break
        if conflict == False:
            self.__rects.append(rect)
            self.__images.append(target_image)
            self.__class_ids.append(class_id)
            return True
        return False

    # 重複率
    def __multiplicity(self, a, b):
        (ax_mn, ay_mn) = a[0]
        (ax_mx, ay_mx) = a[1]
        (bx_mn, by_mn) = b[0]
        (bx_mx, by_mx) = b[1]
        a_area = (ax_mx - ax_mn + 1) * (ay_mx - ay_mn + 1)
        b_area = (bx_mx - bx_mn + 1) * (by_mx - by_mn + 1)
        abx_mn = max(ax_mn, bx_mn)
        aby_mn = max(ay_mn, by_mn)
        abx_mx = min(ax_mx, bx_mx)
        aby_mx = min(ay_mx, by_mx)
        w = max(0, abx_mx - abx_mn + 1)
        h = max(0, aby_mx - aby_mn + 1)
        intersect = w * h
        return intersect / (a_area + b_area - intersect)


# 各クラスのデータ数が同一になるようにカウントする
class Counter:
    def __init__(self, max):
        self.__counter = np.zeros(max)

    def get(self):
        n = np.argmin(self.__counter)
        return int(n)

    def inc(self, index):
        self.__counter[index] += 1

    def print(self):
        print(self.__counter)


def main():
    # 出力先の初期化
    if os.path.exists(OUTPUT_PATH):
        shutil.rmtree(OUTPUT_PATH)
    os.mkdir(OUTPUT_PATH)

    target = Target(TARGET_IMAGE_PATH, BASE_WIDTH, CLASS_NAME)
    background = Background(BACKGROUND_IMAGE_PATH)

    transformer = Transformer(BACK_WIDTH, BACK_HEIGHT)
    manifest = Manifest(CLASS_NAME)
    counter = Counter(len(CLASS_NAME))
    effecter = Effecter()

    no = 0

    while True:
        # 背景画像の取得
        background_image = background.get()

        # 商品データ

        # 重なり率（これを超える場合は、配置されない）
        # rate = 0.1
        rate = 0  # ナンバープレートの重なりは、対象画となるため、重なりを排除する
        data = Data(rate)
        # for _ in range(20):
        for _ in range(10):
            # 現時点で作成数の少ないクラスIDを取得
            class_id = counter.get()
            # 商品画像の取得
            target_image = target.get(class_id)
            # 変換
            (transform_image, rect) = transformer.warp(target_image)
            frame = marge_image(background_image, transform_image)
            # 商品の追加（重複した場合は、失敗する）
            ret = data.append(transform_image, rect, class_id)
            if ret:
                counter.inc(class_id)

        print("max:{}".format(data.max()))
        frame = background_image
        for index in range(data.max()):
            (target_image, _, _) = data.get(index)
            # 合成
            frame = marge_image(frame, target_image)

        # アルファチャンネル削除
        frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)

        # エフェクト
        frame = effecter.gauss(frame, random.randint(0, 2))
        frame = effecter.noise(frame)

        # 画像名
        fileName = "{:05d}.png".format(no)
        no += 1

        # 画像保存
        cv2.imwrite("{}/{}".format(OUTPUT_PATH, fileName), frame)
        # manifest追加
        manifest.appned(fileName, data, frame.shape[0], frame.shape[1])

        for i in range(data.max()):
            (_, rect, class_id) = data.get(i)
            # バウンディングボックス描画（確認用）
            frame = box(frame, rect, class_id)

        counter.print()
        print("no:{}".format(no))
        if MAX <= no:
            break

        # 表示（確認用）
        cv2.imshow("frame", frame)
        cv2.waitKey(1)

    # manifest 保存
    with open("{}/{}".format(OUTPUT_PATH, manifestFile), "w") as f:
        f.write(manifest.get())


main()

"""
Ground Truth形式のデータセットをYolo用に変換する
"""

import json
import glob
import os
import shutil

# 定義
inputPath = "./dataset/output_ground_truth"
# outputPath = "./dataset/yolo"
outputPath = "./yolo"
manifest = "output.manifest"
# 学習用と検証用の分割比率
ratio = 0.8  # 8対2に分割する


# 1件のJデータを表現するクラス
class Data:
    def __init__(self, src):
        # プロジェクト名の取得
        for key in src.keys():
            index = key.rfind("-metadata")
            if index != -1:
                projectName = key[0:index]

        # メタデータの取得
        metadata = src[projectName + "-metadata"]
        class_map = metadata["class-map"]

        # 画像名の取得
        self.imgFileName = os.path.basename(src["source-ref"])
        self.baseName = self.imgFileName.split(".")[0]
        # 画像サイズの取得
        project = src[projectName]
        image_size = project["image_size"]
        self.img_width = image_size[0]["width"]
        self.img_height = image_size[0]["height"]

        self.annotations = []
        # アノテーションの取得
        for annotation in project["annotations"]:
            class_id = annotation["class_id"]
            top = annotation["top"]
            left = annotation["left"]
            width = annotation["width"]
            height = annotation["height"]

            self.annotations.append(
                {
                    "label": class_map[str(class_id)],
                    "width": width,
                    "top": top,
                    "height": height,
                    "left": left,
                }
            )

    # 指定されたラベルを含むかどうか
    def exsists(self, label):
        for annotation in self.annotations:
            if annotation["label"] == label:
                return True
        return False

    def store(self, imagePath, labelPath, inputPath, labels):
        cls_list = []
        for label in labels:
            cls_list.append(label[0])

        text = ""
        for annotation in self.annotations:
            cls_id = cls_list.index(annotation["label"])
            top = annotation["top"]
            left = annotation["left"]
            width = annotation["width"]
            height = annotation["height"]

            yolo_x = (left + width / 2) / self.img_width
            yolo_y = (top + height / 2) / self.img_height
            yolo_w = width / self.img_width
            yolo_h = height / self.img_height

            text += "{} {:.6f} {:.6f} {:.6f} {:.6f}\n".format(
                cls_id, yolo_x, yolo_y, yolo_w, yolo_h
            )
        # txtの保存
        with open("{}/{}.txt".format(labelPath, self.baseName), mode="w") as f:
            f.write(text)
        # 画像のコピー
        shutil.copyfile(
            "{}/{}".format(inputPath, self.imgFileName),
            "{}/{}".format(imagePath, self.imgFileName),
        )


# dataListをラベルを含むものと、含まないものに分割する
def deviedDataList(dataList, label):
    targetList = []
    unTargetList = []
    for data in dataList:
        if data.exsists(label):
            targetList.append(data)
        else:
            unTargetList.append(data)
    return (targetList, unTargetList)


# ラベルの件数の少ない順に並べ替える(配列のインデックスが、クラスIDとなる)
def getLabel(dataList):
    labels = {}
    for data in dataList:
        for annotation in data.annotations:
            label = annotation["label"]
            if label in labels:
                labels[label] += 1
            else:
                labels[label] = 1
    # ラベルの件数の少ない順に並べ替える(配列のインデックスが、クラスIDとなる)
    labels = sorted(labels.items(), key=lambda x: x[1])
    return labels


# 全てのJSONデータを読み込む
def getDataList(inputPath, manifest):
    dataList = []
    with open("{}/{}".format(inputPath, manifest), "r") as f:
        srcList = f.read().split("\n")
        for src in srcList:
            if src != "":
                json_src = json.loads(src)
                dataList.append(Data(json.loads(src)))
    return dataList


def main():
    # 出力先フォルダ生成
    train_images = "{}/train/images".format(outputPath)
    validation_images = "{}/valid/images".format(outputPath)
    train_labels = "{}/train/labels".format(outputPath)
    validation_labels = "{}/valid/labels".format(outputPath)

    os.makedirs(outputPath, exist_ok=True)
    os.makedirs(train_images, exist_ok=True)
    os.makedirs(validation_images, exist_ok=True)
    os.makedirs(train_labels, exist_ok=True)
    os.makedirs(validation_labels, exist_ok=True)

    # 全てのJSONデータを読み込む
    dataList = getDataList(inputPath, manifest)
    log = "全データ: {}件 ".format(len(dataList))

    # ラベルの件数の少ない順に並べ替える(配列のインデックスが、クラスIDとなる)
    labels = getLabel(dataList)
    for i, label in enumerate(labels):
        log += "[{}]{}: {}件 ".format(i, label[0], label[1])
    print(log)

    # 保存済みリスト
    storedList = []
    log = ""
    # ラベルの数の少ないものから優先して分割する
    for i, label in enumerate(labels):
        log = ""
        log += "{} => ".format(label[0])
        # dataListをラベルが含まれるものと、含まないものに分割する
        (targetList, unTargetList) = deviedDataList(dataList, label[0])
        # 保存済みリストから、当該ラベルで既に保存済の件数をカウントする
        (include, notInclude) = deviedDataList(storedList, label[0])
        storedCounst = len(include)
        # train用に必要な件数
        # count = int(label[1] * ratio) - storedCounst
        count = int((len(dataList) * ratio))
        log += "{}:".format(count)
        # train側への保存
        for i in range(count):
            data = targetList.pop()
            data.store(train_images, train_labels, inputPath, labels)
            storedList.append(data)
        # validation側への保存
        log += "{} ".format(len(targetList))
        for data in targetList:
            data.store(validation_images, validation_labels, inputPath, labels)
            storedList.append(data)

        dataList = unTargetList
        log += "残り:{}件".format(len(dataList))
        print(log)


main()

import time
import json
import os
import codecs
import requests
import cv2
import numpy as np
from PIL import ImageFont, ImageDraw, Image


def readApi(im_file_name):
    with open(im_file_name, "rb") as f:
        data = f.read()

    subscription_key = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
    endpoint = "https://japaneast.api.cognitive.microsoft.com/"
    model_version = "2022-04-30"
    # model_version = "2022-01-30-preview"
    # model_version = "2021-04-12"
    language = "ja"

    text_recognition_url = endpoint + "vision/v3.2/read/analyze"
    headers = {
        "Ocp-Apim-Subscription-Key": subscription_key,
        "Content-Type": "application/octet-stream",
    }
    params = {"language ": language, "model-version": model_version}

    response = requests.post(
        text_recognition_url, headers=headers, params=params, json=None, data=data
    )
    response.raise_for_status()

    operation_url = response.headers["Operation-Location"]
    analysis = {}
    poll = True

    while poll:
        response_final = requests.get(
            response.headers["Operation-Location"], headers=headers
        )
        analysis = response_final.json()

        print(json.dumps(analysis, indent=4, ensure_ascii=False))

        time.sleep(1)
        if "analyzeResult" in analysis:
            poll = False
        if "status" in analysis and analysis["status"] == "failed":
            poll = False
    return analysis


def putText(img, text, point, size, color):
    fontFace = "gosic.ttc"
    stroke_width = 1

    # cv2 -> PIL
    imgPIL = Image.fromarray(cv2.cvtColor(img.copy(), cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(imgPIL)
    fontPIL = ImageFont.truetype(font=fontFace, size=size)
    draw.text(xy=point, text=text, fill=color, font=fontPIL, stroke_width=stroke_width)
    # PIL -> cv2
    return cv2.cvtColor(np.array(imgPIL, dtype=np.uint8), cv2.COLOR_RGB2BGR)


def convert_to_gray(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)


def detect(name, image, cls_id, class_name, conf):
    work_dir = "./ocr"
    if not os.path.exists(work_dir):
        os.mkdir(work_dir)

    # 定形のサイズに変換する
    img_all = cv2.resize(image, (440, 220))
    # 業務車の場合は反転する
    if cls_id == 1 or cls_id == 3:
        img_all = cv2.bitwise_not(img_all)
    # グレースケール
    img_all = convert_to_gray(img_all)

    # 部分切り取り
    img_parts = []
    text_parts = ["", "", ""]
    # 上段
    img_parts.append(img_all[0:100, 0:440])
    # ひらがな
    img_parts.append(img_all[100:200, 20:100])
    # 4桁の数字
    img_parts.append(img_all[80:220, 80:440])

    text = ""

    for index in range(len(img_parts)):
        im_file_name = "{}/number_{}_{}.png".format(work_dir, name, index)
        json_file_name = "{}/analysis_{}_{}.json".format(work_dir, name, index)

        if not os.path.isfile(json_file_name):
            # 画像保存
            cv2.imwrite(im_file_name, img_parts[index])
            analysis = readApi(im_file_name)
            with codecs.open(json_file_name, "w+", "utf-8") as fp:
                json.dump(analysis, fp, ensure_ascii=False, indent=2)

        else:
            with open(json_file_name) as f:
                analysis = json.load(f)

        if "analyzeResult" in analysis:
            analyzeResult = analysis["analyzeResult"]

            if "readResults" in analyzeResult:
                readResults = analyzeResult["readResults"]
                readResult = readResults[0]
                if "lines" in readResult:
                    lines = readResult["lines"]
                    for line in lines:
                        x1, y1, _, _, x3, y3, _, _ = line["boundingBox"]
                        img_parts[index] = cv2.rectangle(
                            img_parts[index],
                            (x1, y1),
                            (x3, y3),
                            (0, 0, 255),
                            2,
                            cv2.LINE_4,
                        )
                        # ゴミの排除
                        text = (
                            line["text"]
                            .replace("*", "")
                            .replace("°", "")
                            .replace("-", "")
                        )
                        text_parts[index] = text.replace("*", "")

            # text = text.replace("*", "")

    # 白バックの画像生成
    verify_image = np.zeros((700, 500, 3), dtype=np.uint8) + 255
    h, w = image.shape[:2]

    # 元イメージ
    margin_x = 10
    margin_y = 50
    verify_image[margin_y : h + margin_y, margin_x : w + margin_x, :] = image

    margin_y = 230

    for img in img_parts:
        h2, w2 = img.shape[:2]
        verify_image[margin_y : h2 + margin_y, margin_x : w2 + margin_x, :] = img
        margin_y += h2 + 10

    margin_y = 600
    verify_image = putText(
        verify_image,
        "confidence: {:.2f}".format(conf),
        (margin_x, margin_y),
        15,
        (0, 0, 0),
    )

    margin_y += 25
    verify_image = putText(
        verify_image, class_name, (margin_x, margin_y), 15, (0, 0, 0)
    )
    margin_y += 30
    verify_image = putText(
        verify_image, text_parts[0], (margin_x, margin_y), 25, (0, 0, 0)
    )
    margin_x += 110
    verify_image = putText(
        verify_image, text_parts[1], (margin_x, margin_y), 25, (0, 0, 0)
    )
    margin_x += 50
    verify_image = putText(
        verify_image, text_parts[2], (margin_x, margin_y), 25, (0, 0, 0)
    )

    return verify_image

from ultralytics import YOLO
import cv2
import glob
import os
import ocr

confidence = 0.8
imgsz = 1280
class_names = ["普通自動車（自家用）", "普通自動車（業務用）", "軽自動車（自家用）", "普通自動車（業務用）"]
color = [(0, 0, 255), (255, 0, 255), (255, 0, 0), (255, 255, 0)]
face = ["🐹", "🐰", "🐲", "🐻"]


def main():
    model = YOLO("./best.pt")

    files = glob.glob("./images/*.png")
    for file in files:
        print(file)
        name = os.path.basename(file).split(".")[0]
        img = cv2.imread(file)

        results = model.predict(file, conf=confidence, imgsz=imgsz)

        for result in results:
            boxes = result.boxes.cpu().numpy()
            ids = result.boxes.cls.cpu().numpy().astype(int)
            confs = result.boxes.conf.cpu().numpy()

            if len(boxes) > 0:
                box = boxes[0]
                cls_id = ids[0]
                conf = confs[0]

                # ナンバープレート画像（OCRでの読み取りに使用する）
                r = box.xyxy[0].astype(int)
                number_plate_image = img[r[1] : r[3], r[0] : r[2]]
                verify_image = ocr.detect(
                    name, number_plate_image, cls_id, class_names[cls_id], conf
                )

                img = cv2.rectangle(img, (r[0], r[1]), (r[2], r[3]), color[cls_id], 7)
                print(
                    "\n{} {} conf:{:.2f}  ({},{}), ({},{})".format(
                        face[cls_id], class_names[cls_id], conf, r[0], r[1], r[2], r[3]
                    )
                )

        # 画像の表示
        img = cv2.resize(img, None, fx=0.5, fy=0.5)
        cv2.imshow("image", img)

        cv2.imshow("verify", verify_image)
        cv2.waitKey(1)

        # キー入力待機
        input("Please enter key. ")


if __name__ == "__main__":
    main()