AIFotoOnlus/det.py

###
# Run as following:
# `python3 det.py -d ./images/prova -c newResult.csv`
# __  -d : the directory where you wanna to scan files
# __  -c : the output file name which you are going to make that contain the result of scanning
###


import os
import tempfile
import subprocess
import csv
from imutils.object_detection import non_max_suppression
import numpy as np
import argparse
import time
import cv2
from PIL import Image
import shutil
import pytesseract

ap = argparse.ArgumentParser()
ap.add_argument('-d', '--directory', type=str,
                help="the name of the directory to scan")
ap.add_argument('-c', '--csv', type=str,
                help="the name of the output file which have the result of scanning")
args = vars(ap.parse_args())

directory = args['directory']
outputFile = args['csv']


detecion_net = cv2.dnn.readNet('models/detection.weights', 'models/detection.cfg')
recognition_net = cv2.dnn.readNet('models/recognition.weights', 'models/recognition.cfg')


image_width = 0
image_height = 0
number_classes = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']


def get_output_layers(net):
    layer_names = net.getLayerNames()

    unconnected = net.getUnconnectedOutLayers()
    indices = []
    if isinstance(unconnected, np.ndarray):
        indices = unconnected.flatten()
    else:
        try:
            indices = [u[0] if hasattr(u, '__iter__') else u for u in unconnected]
        except Exception:
            indices = list(unconnected)

    output_layers = [layer_names[int(i) - 1] for i in indices]

    return output_layers


def check_min_value(a):
    if a < 0:
        return 0
    else:
        return a


def check_X_max_value(x):
    if x > image_width:
        return image_width
    else:
        return x


def check_Y_max_value(y):
    if y > image_height:
        return image_height
    else:
        return y


def crop_image(img, x, y, w, h):
    global crop_img_counter

    x1 = check_min_value(x - int(w * 0.1))
    x2 = check_X_max_value(x + w + int(w * 0.1))
    y1 = check_min_value(y - int(h * 0.1))
    y2 = check_Y_max_value(y + h + int(h * 0.1))
    crop_img = img[y1:y2, x1:x2]

    return crop_img


def draw_bounding_box(img, class_id, label, x, y, x_plus_w, y_plus_h):
    # label = str(plate_classes[class_id])

    color = (0, 255, 0)

    cv2.rectangle(img, (x, y), (x_plus_w, y_plus_h), color, 2)

    cv2.putText(img, label, (x - 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 3)


def recog_number(image):
    scale = 0.00392

    # read pre-trained model and config file

    # create input blob
    blob = cv2.dnn.blobFromImage(image, scale, (140, 120), (0, 0, 0), True, crop=False)

    # set input blob for the network
    recognition_net.setInput(blob)

    # run inference through the network
    # and gather predictions from output layers
    outs = recognition_net.forward(get_output_layers(recognition_net))

    # initialization
    class_ids = []
    confidences = []
    boxes = []
    center_X = []
    conf_threshold = 0.5
    nms_threshold = 0.4

    # for each detetion from each output layer
    # get the confidence, class id, bounding box params
    # and ignore weak detections (confidence < 0.5)
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5:
                center_x = int(detection[0] * image.shape[1])
                center_y = int(detection[1] * image.shape[0])
                w = int(detection[2] * image.shape[1])
                h = int(detection[3] * image.shape[0])
                x = center_x - w / 2
                y = center_y - h / 2
                class_ids.append(class_id)
                confidences.append(float(confidence))
                boxes.append([x, y, w, h])
                center_X.append(center_x)

    # apply non-max suppression
    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

    # go through the detections remaining
    # after nms and draw bounding box

    result = ''

    valid_boxes = []
    valid_classids = []
    valid_centerX = []

    for i in indices:
        if isinstance(i, (list, tuple, np.ndarray)):
            idx = int(i[0])
        else:
            idx = int(i)
        box = boxes[idx]
        x = box[0]
        valid_boxes.append(box)
        valid_classids.append(class_ids[idx])
        valid_centerX.append(x)

    for i in range(0, len(valid_centerX)):
        for j in range(i + 1, len(valid_centerX)):
            if valid_centerX[i] > valid_centerX[j]:
                temp = valid_centerX[i]
                valid_centerX[i] = valid_centerX[j]
                valid_centerX[j] = temp

                tem = valid_classids[i]
                valid_classids[i] = valid_classids[j]
                valid_classids[j] = tem

    for i in range(0, len(valid_classids)):

        result += number_classes[valid_classids[i]]


    return result


def recog_text(image):
    scale = 0.00392

    # read pre-trained model and config file

    # create input blob
    blob = cv2.dnn.blobFromImage(image, scale, (416, 416), (0, 0, 0), True, crop=False)

    # set input blob for the network
    detecion_net.setInput(blob)

    # run inference through the network
    # and gather predictions from output layers
    outs = detecion_net.forward(get_output_layers(detecion_net))

    # initialization
    class_ids = []
    confidences = []
    boxes = []
    center_Y_list = []
    conf_threshold = 0.5
    nms_threshold = 0.4

    # for each detetion from each output layer
    # get the confidence, class id, bounding box params
    # and ignore weak detections (confidence < 0.5)
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5:
                center_x = int(detection[0] * image.shape[1])
                center_y = int(detection[1] * image.shape[0])
                w = int(detection[2] * image.shape[1])
                h = int(detection[3] * image.shape[0])
                x = center_x - w / 2
                y = center_y - h / 2
                class_ids.append(class_id)
                confidences.append(float(confidence))
                boxes.append([x, y, w, h])
                center_Y_list.append(center_y)

    # apply non-max suppression
    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

    # go through the detections remaining
    # after nms and draw bounding box

    text = ""

    for i in indices:
        if isinstance(i, (list, tuple, np.ndarray)):
            idx = int(i[0])
        else:
            idx = int(i)
        box = boxes[idx]
        x = box[0]
        y = box[1]
        w = box[2]
        h = box[3]

        plate_img = crop_image(image, round(x), round(y), round(w), round(h))

        license_str = recog_number(plate_img)

        draw_bounding_box(image, class_ids[idx], license_str, round(x), round(y), round(x + w), round(y + h))

        print(license_str)

        text += license_str + ","

    return text[:-1]


class detector:
    def __init__(self):
        super().__init__()
        self.superdir = ""
        self.subdir = directory
        self.outputName = outputFile

        # self.finddir()
        self.ocr()

    def finddir(self):
        currentPath = os.getcwd()
        self.superdir = currentPath + "/images"
        self.subdir = [self.superdir + "/" +
                       sd for sd in os.listdir(self.superdir) if not ("." in sd)]

    def ocr(self):
        if os.path.isfile(self.outputName):
            os.remove(self.outputName)
        for fil in os.listdir(self.subdir):
            if ".jpg" in fil or ".JPG" in fil:
                filepath = self.subdir + "/" + fil
                text = self.text_detect(filepath, fil)
                # i = self.subdir.index(self.subdir)
                # subfolders = [sf for sf in os.listdir(self.superdir) if not ("." in sf)]
                csvPath = self.outputName

                file_exists = os.path.isfile(csvPath)
                if text == "":
                    continue
                with open(csvPath, mode='a') as csv_file:
                    fieldnames = ['filename', 'text']
                    writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
                    if not file_exists:
                        writer.writeheader()
                    writer.writerow({'filename': fil, 'text': text})

    def text_detect(self, filepath, file):
        global image_width, image_height

        image = cv2.imread(filepath)

        image_width = image.shape[1]
        image_height = image.shape[0]

        text = recog_text(image)
        print(text)

        #image = cv2.resize(image, (1000, 800))
        #cv2.imshow("Result", image)
        #cv2.waitKey(0)

        return text


if __name__ == "__main__":
    if directory == None or outputFile == None:
        print('\tInvalid argument inputed. You should input correct arguments.\n\t--e.g.\t `python3 det.py -d ./images/prova -c prova.csv')
    else:
        if not os.path.isdir(directory) or not '.csv' in outputFile:
            if not os.path.isdir(directory):
                print('\tSuch directory doesn\'t exist.')
            else:
                print('\tInvalid argument inputed. You should input correct arguments.\n\t--e.g.\t `python3 det.py -d ./images/prova -c prova.csv')
        else:
            detector()