def writeDatabase(outDBNames,
files_list,
w,
h,
data_root="",
gpu_id=0,
prev_model_basename="",
test_db=False,
debugFlag=False,
trainAfter=False):
"""
Creates two tensors of image matrices. The images are read from the filenames in files_list, resized, converted to grayscale if they are color.
The images are shuffled to avoid statistical issues with caffe, then randomly swapped.
X1 contains just images from each pair, and X2 contains HOG descriptors from the other image of each pair
The end product is created in buffers and written to two LMDBs for X1 and X2 accordingly
"""
if (not data_root == "") and (not data_root.endswith(
'/')): # expect directory name to have '/' appended
data_root += '/'
if not 'linux' in platform:
raise Exception('Error: only UNIX machines are currently supported')
print '\npreparing to transform images and write databases ...\n'
n = len(files_list) # number of samples
n_comp = 3648 # HOG vector length
mem = virtual_memory()
num_buff, bytesNeeded = calcNumBuff(
w, h, n, n_comp, mem) # number of buffers for database writing
if debugFlag:
num_buff = 3
bytesNeeded = 0
n_per_buff = int(math.ceil(n / num_buff))
print bcolors.PURPLE + "Number of buffers: ", num_buff, ", Images per buffer: ", n_per_buff, ", Total image count: " + str(
n) + '\n\n\n\n\n\n\n' + bcolors.ENDC
r = Random(0) # make a random number generator instance seeded with 0
if test_db:
plt.ion()
inds = range(n)
# shuffling indices will slow down the array accessing process,
# but will take away any relationship between adjacent images, making the model better
r.shuffle(inds) # note that shuffle works in place
if not os.path.isdir(data_root + "train_data"):
os.makedirs(data_root + "train_data")
# prepare the max database size. There is no harm in making it too big, since this is only the cap, not the actual size. If disk space runs out, it will throw an error and crash anyways
map_size = 1024**4 # 1 TB
chan = 1
first_buff_flag = True
im_count_tot = 0 # total number of pictures
i_to_show = r.randint(0, n_per_buff - 1)
hog = cv2.HOGDescriptor((16, 32), (16, 16), (16, 16), (8, 8), 2, 1)
X1_db_name = data_root + "train_data/" + outDBNames[0]
X2_db_name = data_root + "train_data/" + outDBNames[1]
with click.progressbar(range(num_buff), label="Total Progress") as bigBar:
for j in bigBar:
##### Database writing #################################
db1 = lmdb.open(X1_db_name, map_size=map_size)
db2 = lmdb.open(X2_db_name, map_size=map_size)
k = 0 # index in X row
txn1 = db1.begin(
write=True,
buffers=True) # start a new transaction for the database
txn2 = db2.begin(write=True, buffers=True)
with click.progressbar(
inds[(j * n_per_buff):((j + 1) * n_per_buff)],
label=("Progress in buffer " + str(j + 1) + " out of " +
str(num_buff))) as bar:
for i in bar: # index in files_list, which is n long
im_file = files_list[i]
im = cv2.imread(im_file)
#only for display
cv2.imshow("a", im)
cv2.waitKey(0)
while im is None: # Some images get corrupted sometimes. Check for this so that it doesnt crash a multi-day running process (sigh)
print "\n\n\nSkipping corrupted image:", im_file, ". Bootstrapping random image from dataset\n\n\n"
im_file = files_list[r.randint(0, n - 1)]
im = cv2.imread(im_file)
if len(im.shape
) > 2: # if color image, convert to grayscale
im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
im = cv2.resize(im, (w, h), interpolation=cv2.INTER_CUBIC)
im_warp = randPerspectiveWarp(
im, w, h, r) # get the perspective warped picture
r.seed(
i
) # adds extra randomness, but is still reproduceable with the same dataset
# image processing function. Needs random number generator
#im, im_warp = preprocess(im, im_warp, r)
if test_db and i == i_to_show: # only show the last image. If we show all of them, computation is VERY slow
plt.subplot(121), plt.imshow(
im, cmap='gray'), plt.title('Original, Brightness')
plt.subplot(122), plt.imshow(
im_warp,
cmap='gray'), plt.title('Perspective Changed')
plt.pause(
10
) # only show it for 10 seconds incase the user walks away.
plt.close()
# randomly choose whether the original or the transformed image go in X1/X2 respectively. This prevents bias in the model since we can only zoom in
switchFlag = r.randint(0, 1)
if switchFlag:
im1 = im_warp
des = hog.compute(cv2.resize(im, (160, 120)))
else:
im1 = im
des = hog.compute(cv2.resize(im_warp, (160, 120)))
str_id = '{:08}'.format(im_count_tot).encode(
'ascii'
) # we only need one str_id as well, since they should be the same for corresponding images
datum1 = caffe.proto.caffe_pb2.Datum()
datum1.channels = chan
datum1.width = w
datum1.height = h
datum1.data = im1.tostring()
txn1.put(str_id,
datum1.SerializeToString()) # add it to database1
datum2 = caffe.proto.caffe_pb2.Datum()
datum2.channels = 1
datum2.width = 1
datum2.height = n_comp
datum2.data = np.reshape(des, (n_comp)).tostring()
txn2.put(str_id,
datum2.SerializeToString()) # add it to database2
k += 1
im_count_tot += 1
if im_count_tot == n - 1:
break
# end for i in inds[(j*n_per_buff):((j+1)*n_per_buff)]:
txn1.commit()
txn2.commit()
# end for j in range(num_buff)
if trainAfter:
if not debugFlag:
batch_sz = 256 # Use 256 batch size for easy testing on small GPU, can increase to train faster. For multi gpu use the bash script with the caffe executable
n_epochs = 43 # We gave our net about 42.25 epochs, so just ceil that number to get 43 epochs
its = int(n_epochs * n /
batch_sz) # calculate number of its based on data size!!
else:
its = 12
snapshot_prefix = '"calc"' # prefix for current model
print "\nDefining net with new database\n"
# define the current train net. Note that the [11:] drops the 'train_data/' prefix
create_net(X1_db_name[11:],
X2_db_name[11:],
max_iter=str(its),
data_root=data_root,
debugFlag=debugFlag)
print "\nInitializing Training"
# train the net with the current data
train("proto/solver.prototxt", GPU_ID=gpu_id)
moveModel(model_dir="calc_" + time.strftime("%d-%m-%Y_%I%M%S")
) # move all the model files to a directory
elif args.subparser=='pass':
from testNet import view_forward_pass
view_forward_pass(args.file1, args.file2)
elif args.subparser=='net':
from makeNet import view_output_size, train, create_net
if args.defineFlag:
# create the net, save the prototext, binary, and solver state
max_iter = args.max_iter
if args.debugFlag:
max_iter = '12'
create_net(args.X1_lmdb, args.X2_lmdb, max_iter=max_iter, debugFlag=args.debugFlag)
if args.trainFlag:
train('proto/solver.prototxt', snapshot_solver_path=args.snapshotProto, init_weights=args.weightsPath)
if args.viewSizeFlag:
view_output_size()
elif args.subparser=='test':
from testNet import plot
plot(net_model_path=args.model_path, data_path=args.data_path, num_include=args.num_include, title=args.title, resize_net=args.rz_net, alexnet_proto_path=args.alex_proto, alexnet_weights=args.alex_weights)
elif args.subparser=='plot':
from testNet import plot_var_db_size
plot_var_db_size(args.calc_fl, args.dbow_fl)