177 lines
6.9 KiB
Python
177 lines
6.9 KiB
Python
import subprocess
|
|
import os
|
|
import shutil
|
|
from time import sleep
|
|
import sys
|
|
from utils import parseOptions
|
|
from utils import getOption
|
|
|
|
parseOptions()
|
|
|
|
DEBUG = getOption("debug")
|
|
WIDTH = str(30)
|
|
HEIGHT = str(30)
|
|
DEFAULT_NEG = 1000
|
|
DEFAULT_POS = 500
|
|
DEFAULT_STAGES = str(5)
|
|
USE_USER_IMG_NUM = False
|
|
|
|
if getOption("width") != None:
|
|
WIDTH = getOption("width")
|
|
if getOption("height") != None:
|
|
HEIGHT = getOption("height")
|
|
if getOption("num_stages") != None:
|
|
DEFAULT_STAGES = getOption("num_stages")
|
|
if getOption("images") != None:
|
|
USE_USER_IMG_NUM = True
|
|
DEFAULT_POS = int(getOption("images"))
|
|
DEFAULT_NEG = DEFAULT_POS * 2
|
|
|
|
GEN_FOLDER = "generated"
|
|
NEGATIVES_FOLDER = "negatives/"
|
|
POSITIVES_FOLDER = "positives/"
|
|
DATA_FOLDER = "haarcascade"
|
|
SAMPLES_FOLDER = "samples"
|
|
VEC_FOLDER = "vectors"
|
|
|
|
NEG_FILE = "negatives.txt"
|
|
POS_FILE = "positives.txt"
|
|
INFO_FILE = "info.lst"
|
|
VEC_FILE = "positives.vec"
|
|
|
|
MERGE_VEC_SCRIPT_PATH = "./utils/mergevec.py"
|
|
|
|
GEN_NORMALIZED_NEGATIVES_FOLDER = "normalized_neg"
|
|
GEN_NORMALIZED_POSITIVES_FOLDER = "normalized_pos"
|
|
|
|
GEN_NEG_FILE = os.path.join(GEN_FOLDER, NEG_FILE)
|
|
GEN_INFO_FILE = os.path.join(GEN_FOLDER, INFO_FILE)
|
|
|
|
GEN_NORMALIZED_NEGATIVES_PATH = os.path.join(GEN_FOLDER, GEN_NORMALIZED_NEGATIVES_FOLDER)
|
|
GEN_NORMALIZED_POSITIVES_PATH = os.path.join(GEN_FOLDER, GEN_NORMALIZED_POSITIVES_FOLDER)
|
|
GEN_SAMPLES_PATH = os.path.join(GEN_FOLDER, SAMPLES_FOLDER)
|
|
GEN_VEC_PATH = os.path.join(GEN_FOLDER, VEC_FOLDER)
|
|
|
|
CREATE_SAMPLES_COMMAND_EX = "opencv_createsamples -img POS_IMG -bg " + GEN_NEG_FILE + " -info INFO_LOC -pngoutput SAMPLES_PATH -maxxangle 0.5 -maxyangle 0.5 -maxzangle 0.5 -num NUM_IMG"
|
|
CREATE_NEG_FILE_EX = 'find ' + GEN_NORMALIZED_NEGATIVES_FOLDER + ' -iname "*.jpg" > ' + NEG_FILE
|
|
CREATE_POS_FILE_EX = 'find ' + GEN_NORMALIZED_POSITIVES_FOLDER + ' -iname "*.jpg" > POS_FILE'
|
|
CREATE_VEC_COMMAND_EX = "opencv_createsamples -info INFO_LOC -num NUM_IMG -w " + WIDTH + " -h " + HEIGHT + " -vec VEC_FILE"
|
|
TRAIN_CASCADE_COMMAND = "opencv_traincascade -data ../" + DATA_FOLDER + " -vec " + VEC_FILE + " -bg " + NEG_FILE + " -numPos POS_IMGS -numNeg NEG_IMGS -numStages " + DEFAULT_STAGES + " -w " + WIDTH + " -h " + HEIGHT + " -featureType HAAR -precalcValBufSize 2048 -precalcIdxBufSize 2048"
|
|
MERGE_VECTORS_COMMAND_EX = "python " + MERGE_VEC_SCRIPT_PATH + " -v INPUT_FILES -o OUTPUT_FILE"
|
|
|
|
if __name__ == '__main__':
|
|
# Ensuring the previous build has been removed
|
|
if os.path.exists(GEN_FOLDER):
|
|
if DEBUG:
|
|
print("[*] Removing previous build...")
|
|
shutil.rmtree(GEN_FOLDER)
|
|
if DEBUG:
|
|
print("[*] Previous build has been removed")
|
|
sleep(0.1)
|
|
|
|
# Making required folders
|
|
if DEBUG:
|
|
print("[*] Beginning folder creation...")
|
|
subprocess.Popen(["mkdir", GEN_FOLDER]).wait()
|
|
subprocess.Popen(["mkdir", GEN_NORMALIZED_NEGATIVES_PATH]).wait()
|
|
subprocess.Popen(["mkdir", GEN_NORMALIZED_POSITIVES_PATH]).wait()
|
|
subprocess.Popen(["mkdir", GEN_VEC_PATH]).wait()
|
|
if DEBUG:
|
|
print("[*] Completed folder creation")
|
|
sleep(0.1)
|
|
|
|
# Copying all negative images into a new folder
|
|
if DEBUG:
|
|
print("[*] Copying all negative images to secondary location...")
|
|
NUM_NEGATIVES = 0
|
|
for file_name in os.listdir(NEGATIVES_FOLDER):
|
|
full_file_name = os.path.join(NEGATIVES_FOLDER, file_name)
|
|
|
|
if os.path.isfile(full_file_name):
|
|
if ".jpg" in full_file_name:
|
|
shutil.copy(full_file_name, GEN_NORMALIZED_NEGATIVES_PATH)
|
|
|
|
NUM_NEGATIVES += 1
|
|
if DEBUG:
|
|
print("[*] Completed negative image copying")
|
|
sleep(0.1)
|
|
|
|
# Generate negative images file (contains all of the negative images)
|
|
if DEBUG:
|
|
print("[*] Generating negative images information file...")
|
|
os.chdir(GEN_FOLDER)
|
|
os.system(CREATE_NEG_FILE_EX)
|
|
os.chdir("../")
|
|
if DEBUG:
|
|
print("[*] Completed negative images information generation")
|
|
sleep(0.1)
|
|
|
|
# Copying all positive images into a new folder
|
|
if DEBUG:
|
|
print("[*] Copying all positive images to secondary location...")
|
|
NUM_POSITIVES = 0
|
|
for file_name in os.listdir(POSITIVES_FOLDER):
|
|
full_file_name = os.path.join(POSITIVES_FOLDER, file_name)
|
|
|
|
if os.path.isfile(full_file_name):
|
|
if ".jpg" in full_file_name:
|
|
shutil.copy(full_file_name, GEN_NORMALIZED_POSITIVES_PATH)
|
|
|
|
NUM_POSITIVES += 1
|
|
if DEBUG:
|
|
print("[*] Completed positive image copying")
|
|
sleep(0.1)
|
|
|
|
# Setting the number of images the train_cascade command should use
|
|
if USE_USER_IMG_NUM == False:
|
|
DEFAULT_NEG = int(NUM_NEGATIVES * 0.9)
|
|
DEFAULT_POS = DEFAULT_NEG / 2
|
|
|
|
# Creating samples for the positive images
|
|
counter = 1
|
|
for file_name in os.listdir(GEN_NORMALIZED_POSITIVES_PATH):
|
|
full_file_name = os.path.join(GEN_NORMALIZED_POSITIVES_PATH, file_name)
|
|
|
|
if os.path.isfile(full_file_name):
|
|
if ".jpg" in full_file_name:
|
|
if DEBUG:
|
|
print("[*] Creating samples...")
|
|
subprocess.Popen(["mkdir", (GEN_SAMPLES_PATH + str(counter))]).wait()
|
|
create_samples = CREATE_SAMPLES_COMMAND_EX.replace("POS_IMG", full_file_name).replace("NUM_IMG", str(DEFAULT_POS / 10)).replace("SAMPLES_PATH", (GEN_SAMPLES_PATH + str(counter))).replace("INFO_LOC", os.path.join((GEN_SAMPLES_PATH + str(counter)), INFO_FILE))
|
|
process = subprocess.Popen(create_samples.split(" "), stdout=sys.stdout)
|
|
process.wait()
|
|
if DEBUG:
|
|
print("[*] Finished creating samples")
|
|
|
|
# Creating vector file
|
|
if DEBUG:
|
|
print("[*] Creating vector file...")
|
|
GEN_VEC_FILE = os.path.join(GEN_VEC_PATH, VEC_FILE)
|
|
create_vec_file = CREATE_VEC_COMMAND_EX.replace("NUM_IMG", str(DEFAULT_POS / 10)).replace("VEC_FILE", (GEN_VEC_FILE[:-4] + str(counter) + GEN_VEC_FILE[-4:])).replace("INFO_LOC", os.path.join((GEN_SAMPLES_PATH + str(counter)), INFO_FILE))
|
|
os.system(create_vec_file)
|
|
if DEBUG:
|
|
print("[*] Finished writing vector file")
|
|
|
|
counter += 1
|
|
sleep(0.1)
|
|
|
|
# Merging vector files into one
|
|
if DEBUG:
|
|
print("[*] Compiling vector files into one single file...")
|
|
merge_command = MERGE_VECTORS_COMMAND_EX.replace("INPUT_FILES", GEN_VEC_PATH).replace("OUTPUT_FILE", os.path.join(GEN_FOLDER, VEC_FILE))
|
|
os.system(merge_command)
|
|
if DEBUG:
|
|
print("[*] Finished compiling vector files into a single file")
|
|
sleep(0.1)
|
|
|
|
# Training cascade
|
|
if DEBUG:
|
|
print("[*] Training the cascade...")
|
|
os.chdir(GEN_FOLDER)
|
|
process = subprocess.Popen(TRAIN_CASCADE_COMMAND.replace("POS_IMGS", str(DEFAULT_POS)).replace("NEG_IMGS", str(DEFAULT_NEG)), shell=True)
|
|
process.wait()
|
|
os.chdir("../")
|
|
if DEBUG:
|
|
print("[*] Finished training the cascade")
|
|
sleep(0.1)
|