First commit

det3.py

@@ -0,0 +1,334 @@
+###
+# 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
+from multiprocessing import Pool
+
+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")
+ap.add_argument('-tn', '--textnegative', action='store_true',
+                help="skip files that begin with 'tn_'")
+ap.add_argument('-r', '--recursive', type=int, 
+                help="recursively called with defined concurrent operations")
+args = vars(ap.parse_args())
+
+directory = args['directory']
+outputFile = args['csv']
+textnegative = args['textnegative']
+recursive = args['recursive']
+
+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()
+
+    output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+
+    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:
+        i = i[0]
+        box = boxes[i]
+        x = box[0]
+        valid_boxes.append(box)
+        valid_classids.append(class_ids[i])
+        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:
+        i = i[0]
+        box = boxes[i]
+        x = box[0]
+        y = box[1]
+        w = box[2]
+        h = box[3]
+
+        # if center_Y_list[i] in range(int(image.shape[0] * 0.2), int(image.shape[0] * 0.8)):
+
+        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[i], 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, subdir=None, outputName=None):
+        super().__init__()
+        self.superdir = ""
+        self.subdir = subdir if subdir else directory
+        self.outputName = outputName if outputName else outputFile
+        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)
+        content = os.listdir(self.subdir)
+        files = [fil for fil in content if (".jpg" in fil or ".JPG" in fil) and os.path.isfile(os.path.join(self.subdir,fil))]
+        if recursive:
+            with Pool(processes=recursive) as pool:
+                pool.map(self.process_file, files)
+        else:
+            for fil in files:
+                self.process_file(fil)
+                
+        # Check for the -r or --recursive argument
+        if recursive:
+            directories = [item for item in content if os.path.isdir(os.path.join(self.subdir, item))]
+            for directory in directories:
+                # Create a new instance of detector for each subdirectory
+                detector(subdir=os.path.join(self.subdir, directory), outputName=self.outputName)
+    
+    def process_file(self, fil):
+        if textnegative and fil.startswith('tn_'):
+            return
+        filepath = self.subdir + "/" + fil
+        text = self.text_detect(filepath, fil)
+        csvPath = self.outputName
+
+        file_exists = os.path.isfile(csvPath)
+        if text == "":
+            return
+        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()