annotate.py

import glob
import os
import cv2
import pytesseract
import utils
import json
import matplotlib.pyplot as plt

ANNOT_FILE = "annotations.json"
DATA_DIR = "/home/prime/Dataset/lines"


if os.path.exists(ANNOT_FILE):
    with open(ANNOT_FILE) as handle:
        output = json.loads(handle.read())
else:
    output = dict()

line_images = glob.glob(os.path.join(DATA_DIR, "*"))


for out in output.keys():
    f = os.path.join(DATA_DIR, out)
    line_images.remove(f)

os.system("clear")
print(f"\nDone -> {len(output.keys())}    Remaining -> {len(line_images)}    Total -> {len(output.keys()) + len(line_images)}\n")


def read_tesseract(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    gray = cv2.resize(gray, None, fx=3, fy=3, interpolation=cv2.INTER_CUBIC)

    ret, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY_INV)

    rect_kern = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
    dilation = cv2.dilate(thresh, rect_kern, iterations=1)

    dilation = cv2.bitwise_not(dilation)

    config = '-c tessedit_char_whitelist=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ --psm 13 --oem 1 --tessdata-dir ./tessdata'
    text = pytesseract.image_to_string(dilation, lang='foo', config=config)
    #config = '-c tessedit_char_whitelist=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ --psm 13 --oem 1'
    #text = pytesseract.image_to_string(dilation, config=config)
    text = utils.remove_alnum(text)

    return gray, text


try:

    for line in line_images:

        image = cv2.imread(line)
        gray, tess_text = read_tesseract(image)

        plt.imshow(gray)
        plt.show(block=False)

        print(f"\033[92m{os.path.basename(line)}\033[0m")
        print("T->",tess_text)
        user_text = input("U-> ")

        if user_text == "":
            label = tess_text.upper()
        else:
            label = user_text.upper()

        output[os.path.basename(line)] = (label, label == tess_text.upper())

        print("Label ->", label)
        print("\n\n\n")

        plt.close()

except KeyboardInterrupt:

    with open(ANNOT_FILE, "w") as handle:
        json.dump(output, handle, indent=4)

    print("\nSaving...")


with open(ANNOT_FILE, "w") as handle:
    json.dump(output, handle, indent=4)

print("\nSaving...")

combine_folders.py

import glob
import os
import shutil

DATASET_FOLDER = "/home/prime/Dataset"
COMPLETE_DATASET = "complete_data"
counter = 1


if os.path.exists(COMPLETE_DATASET):
    print("Complete Dataset folder already exists")
    exit()
else:
    os.mkdir(COMPLETE_DATASET)



folders = glob.glob(os.path.join(DATASET_FOLDER, "*"))

for fold in folders:
    files = glob.glob(os.path.join(fold, "*"))

    for f in files:
        s = os.path.join(COMPLETE_DATASET, f"{counter}.jpeg")
        shutil.copy(f, s)
        print(f, "-", s)
        counter += 1

        

correct_skew.py

import glob
import os
import cv2

plates = sorted(glob.glob("/home/prime/Dataset/plates/*"))
save_dir = "plates_no_skew/"

if os.path.exists(save_dir):
    print("plates_no_skew/ directory exists")
    exit()
else:
    os.mkdir(save_dir)

for plate in plates:
    image = cv2.imread(plate)
    save_path = os.path.join(save_dir, os.path.basename(plate))
    angle, rotated = utils.correct_skew(image)
    cv2.imwrite(save_path, rotated)
    print('.')

encode.py

class Encoder:
  def __init__(self):
    self.char_to_num = {}
    self.num_to_char = {}
    self.populate()
    
  def populate(self):
    self.char_to_num[" "] = -1
    for val in range(0, 10):
      self.char_to_num[str(val)] = val
    number = 10
    for val in range(65, 91):
      self.char_to_num[chr(val)] = number
      number += 1
    
    for key, val in self.char_to_num.items():
      self.num_to_char[val] = key
    
  def encode(self, sentence, max_length):
    output = []
    for s in sentence:
      output.append(self.char_to_num[s])
    
    while len(output) < max_length:
      output.append(self.char_to_num[' '])
    
    return output
  
  def decode(self, sentence):
    output = ""
    for s in sentence:
      output += self.num_to_char[s]
    
    return output.strip()

image_to_plate.py

import glob
import os
import cv2

images = sorted(glob.glob("/home/prime/Dataset/complete_data/*"))
save_dir = "plates/"

if os.path.exists(save_dir):
    print("plates/ directory exists")
    exit()
else:
    os.mkdir("plates/")


CONFIDENCE_THRESHOLD = 0.2
NMS_THRESHOLD = 0.4

net = cv2.dnn.readNet("yolov4.weights", "yolov4.cfg")
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)

model = cv2.dnn_DetectionModel(net)
model.setInputParams(size=(608, 608), scale=1/255, swapRB=True)

for path in images:
    image = cv2.imread(path)
    classes, scores, boxes = model.detect(image, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)

    if len(classes) != 1:
        print(path, "-", len(classes))
    else:
        print(".")

    save_path = os.path.join(save_dir, os.path.basename(path))

    for idx, (classid, score, box) in enumerate(zip(classes, scores, boxes)):
        x, y, w, h = box
        roi_box = image[y: y+h, x: x+w]
        cv2.imwrite(save_path, roi_box)
    

plate_to_line.py

import glob
import os
import cv2

plates = sorted(glob.glob("/home/prime/Dataset/plates_no_skew/*"))
save_dir = "lines/"
counter = 1

if os.path.exists(save_dir):
    print("lines/ directory exists")
    exit()
else:
    os.mkdir(save_dir)


CONFIDENCE_THRESHOLD = 0.2
NMS_THRESHOLD = 0.4

net = cv2.dnn.readNet("yolo_line.weights", "yolo_line.cfg")
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)

model = cv2.dnn_DetectionModel(net)
model.setInputParams(size=(128, 128), scale=1/255, swapRB=True)

for path in plates:

    image = cv2.imread(path)
    classes, scores, boxes = model.detect(image, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)

    for idx, (classid, score, box) in enumerate(zip(classes, scores, boxes)):

        x, y, w, h = box
        roi_box = image[y: y+h, x: x+w]

        save_path = os.path.join(save_dir, f"{counter}.jpeg")
        
        cv2.imwrite(save_path, roi_box)

        counter += 1
        print(".")