92 lines
2.4 KiB
Python
92 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import sys
|
|
import os
|
|
import time
|
|
import glob
|
|
import argparse
|
|
|
|
# input arguments
|
|
parser = argparse.ArgumentParser(description='Prepar dataset for Cascade trainer.\n Defaults: posWidthX=50, posWidthY=50, negWidthX=100, negWidthY=100, location /pos/*.*, /neg/*.*')
|
|
parser.add_argument('-posWidthX', type=int, help='Positive sample x Width', default=80)
|
|
parser.add_argument('-posWidthY', type=int, help='Positive sample y Width', default=60)
|
|
|
|
parser.add_argument('-negWidthX', type=int, help='Negative sample x Width', default=160)
|
|
parser.add_argument('-negWidthY', type=int, help='Negative sample y Width', default=120)
|
|
|
|
parser.add_argument('-pos', type=str, help='Positive samples location', default='/pos/*.*')
|
|
parser.add_argument('-neg', type=str, help='Negative samples location', default='/neg/*.*')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# sample sizes
|
|
POS_SIZE=(args.posWidthX,args.posWidthY)
|
|
NEG_SIZE = (args.negWidthX,args.negWidthY)
|
|
|
|
# current path
|
|
path, filename = os.path.split(os.path.realpath(__file__))
|
|
|
|
# read files
|
|
pfiles = glob.glob(path + args.pos)
|
|
nfiles = glob.glob(path + args.neg)
|
|
|
|
|
|
# check directory structure
|
|
if len(pfiles)== 0:
|
|
print ('no positive images found, check pos dir!')
|
|
exit
|
|
if len(nfiles)== 0:
|
|
print ('no negative images found, check neg dir!')
|
|
exit
|
|
|
|
# create positive images descriptor
|
|
f = open(path +'/info.dat','w')
|
|
|
|
for pf in pfiles:
|
|
|
|
# create tag to write in info.dat => objnr=1 startx=0 starty=0 endx=50 endy=50
|
|
infoline = pf + ' 1 0 0 ' + str(POS_SIZE[0]) + ' ' + str(POS_SIZE[1]) +'\n'
|
|
|
|
f.write(infoline)
|
|
|
|
# resize img to desired size
|
|
img = cv2.imread(pf)
|
|
h,w = img.shape[:2]
|
|
|
|
# resize if picture dimensiona are different
|
|
if (w,h) != POS_SIZE:
|
|
img = cv2.resize(img,POS_SIZE)
|
|
try:
|
|
img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
|
|
img = cv2.equalizeHist(img)
|
|
except:
|
|
pass
|
|
cv2.imwrite(pf,img)
|
|
|
|
f.close()
|
|
|
|
# create negative images descriptor
|
|
f = open(path +'/neg.txt','w')
|
|
|
|
for nf in nfiles:
|
|
# print(nf)
|
|
|
|
f.write(nf + '\n')
|
|
|
|
img = cv2.imread(nf)
|
|
h,w = img.shape[:2]
|
|
|
|
# resize if picture dimensiona are different
|
|
if (w,h) != NEG_SIZE:
|
|
img = cv2.resize(img,NEG_SIZE)
|
|
try:
|
|
img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
|
|
img = cv2.equalizeHist(img)
|
|
except:
|
|
pass
|
|
|
|
cv2.imwrite(nf,img)
|
|
f.close()
|