327 lines
9.8 KiB
Python
327 lines
9.8 KiB
Python
###
|
|
# 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
|
|
# __ -tn : Skip files that begin with 'tn_'
|
|
###
|
|
|
|
|
|
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")
|
|
ap.add_argument('-tn', '--textnegative', action='store_true',
|
|
help="skip files that begin with 'tn_'")
|
|
args = vars(ap.parse_args())
|
|
|
|
directory = args['directory']
|
|
outputFile = args['csv']
|
|
textnegative = args['textnegative']
|
|
|
|
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)
|
|
if license_str: # Only append if license_str is not empty
|
|
text += license_str + ","
|
|
|
|
return text[:-1] if text else text
|
|
|
|
|
|
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 textnegative and fil.startswith('tn_'):
|
|
continue
|
|
if ".jpg" in fil or ".JPG" in fil or ".jpeg" in fil:
|
|
#filepath = self.subdir + "/" + fil
|
|
filepath = os.path.join(self.subdir, fil) # Use os.path.join for cross-platform compatibility
|
|
try:
|
|
text = self.text_detect(filepath, fil)
|
|
except Exception as e:
|
|
print(f"\033[91mError processing {fil}: {e}\033[0m") # Print error in red
|
|
continue # Skip to the next image
|
|
# 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 == "" or text == ",":
|
|
print(f"{fil}: N/A")
|
|
continue
|
|
print(f"{fil}: \033[92m{text}\033[0m")
|
|
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()
|