def inference(modelFilePath: str):
imageSize = 224
sess_options = rt.SessionOptions()
sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
sess = rt.InferenceSession(modelFilePath, sess_options)
input_name = sess.get_inputs()[0].name
print("input_name {}".format(input_name))
mask_name = sess.get_outputs()[0].name
print("mask_name {}".format(mask_name))
totalTime = Timer('total time', 10)
readFrameTime = Timer('read frame', 10)
inferenceTime = Timer('inference', 10)
camera = cv2.VideoCapture(0)
while True:
totalTime.start()
readFrameTime.start()
ret_val, img = camera.read()
readFrameTime.end()
inferenceTime.start()
data = {
'image': img
}
augData = augmentations.valid_transforms(imageSize)(**data)
imgs = [augData['image']]
imgs = np.stack(imgs)
masks = sess.run(None, {input_name: imgs})[0]
inferenceTime.end()
for i in range(masks.shape[0]):
image = imgs[i]
mask = masks[i]
threshold = np.copy(mask)
maskThreshold = 0.3 # 150 / 255
binaryMask = mask[:, :, :] < maskThreshold
threshold[binaryMask] = 0.0
threshold[~binaryMask] = 1.0
ksize = (3, 3)
threshold = cv2.blur(threshold, ksize)
imageMask = np.stack((threshold, threshold, threshold), axis=-1)
background = np.ones_like(image) / 2
cv2.imshow('image', image * imageMask + background * (1.0 - imageMask))
cv2.imshow('mask', mask)
cv2.imshow('threshold', threshold)
cv2.waitKey(1)
totalTime.end()
totalTime.logAverageTime()
readFrameTime.logAverageTime()
inferenceTime.logAverageTime()
cv2.destroyAllWindows()
def main():
args = parse_args()
imagesDir = args.imagesDir
datasetDir = args.datasetDir
modelFilePath = args.modelFilePath
model = tf.keras.models.load_model(
modelFilePath,
compile=False
)
imageSize = 224
imageFilenames = [filename for filename in os.listdir(imagesDir) if os.path.isfile(os.path.join(imagesDir, filename))]
i = 0
for imageFilename in imageFilenames:
ext = os.path.splitext(imageFilename)[1]
if ext != '.jpg':
continue
imagePath = os.path.join(imagesDir, imageFilename)
print("imagePath {}".format(imagePath))
sourceImage = cv2.imread(imagePath)
resized = cv2.resize(sourceImage, (imageSize, imageSize))
data = {
'image': sourceImage
}
augData = augmentations.valid_transforms(imageSize)(**data)
imgs = [augData['image']]
imgs = np.stack(imgs)
image = imgs[0]
masks = model.predict(imgs)
mask = masks[0]
threshold = np.copy(mask)
maskThreshold = 0.3 # 150 / 255
binaryMask = mask[:, :, :] < maskThreshold
threshold[binaryMask] = 0.0
threshold[~binaryMask] = 1.0
ksize = (3, 3)
blurredThreshold = cv2.blur(threshold, ksize)
imageMask = np.stack((blurredThreshold, blurredThreshold, blurredThreshold), axis=-1)
background = np.ones_like(resized) / 2
shutil.copyfile(imagePath, os.path.join(datasetDir, "image{}.jpg".format(i)))
cv2.imwrite(os.path.join(datasetDir, "infMask{}.png".format(i)), threshold)
cv2.imshow('image', resized)
cv2.imshow('masked image', resized / 255 * imageMask + background * (1.0 - imageMask))
cv2.imshow('mask', mask)
cv2.waitKey(1)
i += 1
cv2.destroyAllWindows()
def train(model,
preprocess_input,
humanDataset: SegmentationDataset,
nonHumanDataset: SegmentationDataset,
valHumanDataset: SegmentationDataset,
valNonHumanDataset: SegmentationDataset,
trainingDir: str,
modelEncoder: str,
batchSize: int = 1,
epochs: int = 1,
startEpoch: int = 0,
imageSize=224):
validationPacketSize = 16 * 16
x_val_h, y_val_h = valHumanDataset.readBatch(validationPacketSize)
augmentations.applyTransforms(x_val_h, y_val_h,
augmentations.valid_transforms(imageSize))
x_val_nh, y_val_nh = valNonHumanDataset.readBatch(validationPacketSize)
augmentations.applyTransforms(x_val_nh, y_val_nh,
augmentations.valid_transforms(imageSize))
x_val = np.stack(x_val_h + x_val_nh)
y_val = np.stack(y_val_h + y_val_nh)
x_val = preprocess_input(x_val)
checkPointPath = os.path.join(trainingDir,
'u-net-{}.chpt'.format(modelEncoder))
checkPointCallback = tf.keras.callbacks.ModelCheckpoint(
filepath=checkPointPath, verbose=1)
packetSize = 8 * 8
humanDatasetSequence = HumanDatasetSequence(humanDataset, nonHumanDataset,
packetSize, preprocess_input)
model.fit(x=humanDatasetSequence,
batch_size=batchSize,
shuffle=False,
epochs=epochs,
initial_epoch=startEpoch,
validation_data=(x_val, y_val),
callbacks=[checkPointCallback])
def inference(model, imageSize = 224):
totalTime = Timer('total time', 10)
readFrameTime = Timer('read frame', 10)
inferenceTime = Timer('inference', 10)
camera = cv2.VideoCapture(0)
while True:
totalTime.start()
readFrameTime.start()
ret_val, img = camera.read()
readFrameTime.end()
inferenceTime.start()
data = {
'image': img
}
augData = augmentations.valid_transforms(imageSize)(**data)
imgs = [augData['image']]
imgs = np.stack(imgs)
masks = model.predict(imgs)
inferenceTime.end()
for i in range(masks.shape[0]):
image = imgs[i]
mask = masks[i]
threshold = np.copy(mask)
maskThreshold = 0.3 # 150 / 255
binaryMask = mask[:, :, :] < maskThreshold
threshold[binaryMask] = 0.0
threshold[~binaryMask] = 1.0
ksize = (3, 3)
threshold = cv2.blur(threshold, ksize)
imageMask = np.stack((threshold, threshold, threshold), axis=-1)
background = np.ones_like(image) / 2
cv2.imshow('image', image * imageMask + background * (1.0 - imageMask))
cv2.imshow('mask', mask)
cv2.imshow('threshold', threshold)
cv2.waitKey(1)
totalTime.end()
totalTime.logAverageTime()
readFrameTime.logAverageTime()
inferenceTime.logAverageTime()
cv2.destroyAllWindows()
def explicitTrain(model,
preprocess_input,
humanDataset: SegmentationDataset,
nonHumanDataset: SegmentationDataset,
valHumanDataset: SegmentationDataset,
valNonHumanDataset: SegmentationDataset,
trainingDir: str,
modelEncoder: str,
imageSize=224,
batchSize: int = 1,
epochs: int = 1,
startEpoch: int = 0):
validationPacketSize = 16 * 16
x_val_h, y_val_h = valHumanDataset.readBatch(validationPacketSize)
augmentations.applyTransforms(x_val_h, y_val_h,
augmentations.valid_transforms(imageSize))
x_val_nh, y_val_nh = valNonHumanDataset.readBatch(validationPacketSize)
augmentations.applyTransforms(x_val_nh, y_val_nh,
augmentations.valid_transforms(imageSize))
x_val = np.stack(x_val_h + x_val_nh)
y_val = np.stack(y_val_h + y_val_nh)
x_val = preprocess_input(x_val)
# checkPointPath = os.path.join(trainingDir, 'u-net-{}.chpt'.format(modelEncoder))
# checkPointCallback = tf.keras.callbacks.ModelCheckpoint(filepath=checkPointPath,
# save_weights_only=True,
# verbose=1)
SAVE_AFTER_NUMBER = 50000
packetSize = 16 * 16
nonHumanPacketSize = max(
(packetSize * len(nonHumanDataset)) // len(humanDataset), 1)
csv_logger = CSVLogger('training.log', append=True)
for epoch in range(startEpoch, epochs):
logger.info('epoch %d', epoch)
humanDataset.reset()
nonHumanDataset.reset()
try:
packets = len(humanDataset) // packetSize
for packetIndex in range(packets - 1):
logger.debug('reading batch, memory used %f', usedMemory())
x_train_h, y_train_h = humanDataset.readBatch(packetSize)
x_train_h, y_train_h = augmentations.appendTransforms(
x_train_h, y_train_h,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
logger.debug('reading human batch, memory used %f',
usedMemory())
x_train_nh, y_train_nh = nonHumanDataset.readBatch(
nonHumanPacketSize)
x_train_nh, y_train_nh = augmentations.appendTransforms(
x_train_nh, y_train_nh,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
logger.debug('reading nonHuman batch, memory used %f',
usedMemory())
x_train, y_train = HumanDatasetSequence.shuffleHumanNonHuman(
x_train_h, x_train_nh, y_train_h, y_train_nh)
x_train = np.concatenate((x_train, ))
y_train = np.concatenate((y_train, ))
del x_train_h
del x_train_nh
del y_train_h
del y_train_nh
logger.debug('concatenate batches, memory used %f',
usedMemory())
x_train = preprocess_input(x_train)
# x_train = x_train / 255
logger.debug('preprocess x_train, memory used %f',
usedMemory())
logger.debug('start train on %d samples, memory used %f',
len(x_train), usedMemory())
model.fit(x=x_train,
y=y_train,
batch_size=batchSize,
epochs=epoch + 1,
initial_epoch=epoch,
validation_data=(x_val, y_val),
callbacks=[csv_logger])
saveModel = (
(humanDataset.index + nonHumanDataset.index) %
SAVE_AFTER_NUMBER) < (packetSize + nonHumanPacketSize)
if saveModel:
save_model(model, trainingDir, modelEncoder, packetIndex)
del x_train
del y_train
logger.debug('trained on %d samples, memory used %f',
humanDataset.index + nonHumanDataset.index,
usedMemory())
# gc.collect()
# objgraph.show_most_common_types(limit=50)
# obj = objgraph.by_type('list')[1000]
# objgraph.show_backrefs(obj, max_depth=10)
x_train_h, y_train_h = humanDataset.readBatch(packetSize)
x_train_h, y_train_h = augmentations.appendTransforms(
x_train_h, y_train_h,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
x_train_nh, y_train_nh = nonHumanDataset.readBatch(
nonHumanPacketSize)
x_train_nh, y_train_nh = augmentations.appendTransforms(
x_train_nh, y_train_nh,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
x_train, y_train = HumanDatasetSequence.shuffleHumanNonHuman(
x_train_h, x_train_nh, y_train_h, y_train_nh)
x_train = np.concatenate((x_train, ))
y_train = np.concatenate((y_train, ))
del x_train_h
del x_train_nh
del y_train_h
del y_train_nh
x_train = preprocess_input(x_train)
# x_train = x_train / 255
model.fit(x=x_train,
y=y_train,
batch_size=batchSize,
epochs=1,
validation_data=(x_val, y_val),
callbacks=[csv_logger])
save_model(model, trainingDir, modelEncoder, packets - 1)
del x_train
del y_train
gc.collect()
logger.info('epoch %d is trained', epoch)
except Exception as e:
logger.error('Exception %s', str(e))
traceback.print_exc()
return
now = datetime.now()
dt_string = now.strftime("%Y/%m/%d %H:%M:%S")
modelPath = os.path.join(
trainingDir,
'u-net-{}_epoch{}_{}.tfmodel'.format(modelEncoder, epoch,
dt_string))
model.save(modelPath)
logger.info('model saved')
def explicitTrain(model,
preprocess_input,
humanDataset: SegmentationDataset,
activeLearningDataset: SegmentationDataset,
nonHumanDataset: SegmentationDataset,
valHumanDataset: SegmentationDataset,
valNonHumanDataset: SegmentationDataset,
trainingDir: str,
modelEncoder: str,
batchSize: int = 1,
epochs: int = 1,
startEpoch: int = 0,
imageSize: int = 224):
validationPacketSize = 16 * 16
x_val_h, y_val_h = valHumanDataset.readBatch(validationPacketSize)
augmentations.applyTransforms(x_val_h, y_val_h,
augmentations.valid_transforms(imageSize))
x_val_nh, y_val_nh = valNonHumanDataset.readBatch(validationPacketSize)
augmentations.applyTransforms(x_val_nh, y_val_nh,
augmentations.valid_transforms(imageSize))
x_val = np.stack(x_val_h + x_val_nh)
y_val = np.stack(y_val_h + y_val_nh)
x_val = preprocess_input(x_val)
# checkPointPath = os.path.join(trainingDir, 'u-net-{}.chpt'.format(modelEncoder))
# checkPointCallback = tf.keras.callbacks.ModelCheckpoint(filepath=checkPointPath,
# save_weights_only=True,
# verbose=1)
packetSize = 16 * 16
nonHumanPacketSize = max(
(packetSize * len(nonHumanDataset)) // len(humanDataset), 1)
activeLearningRepeatCount = 3
activeLearningImagesPerPacket = activeLearningRepeatCount * packetSize * len(
activeLearningDataset) / len(humanDataset)
activeLearningImagesRead = 0
activeLearningImagesNeedToRead = 0
print("activeLearningImagesPerPacket {}".format(
activeLearningImagesPerPacket))
for epoch in range(startEpoch, epochs):
logger.info('epoch %d', epoch)
humanDataset.reset()
nonHumanDataset.reset()
try:
packets = len(humanDataset) // packetSize
for packetIndex in range(packets - 1):
logger.debug('reading batch, memory used %f', usedMemory())
x_train_h, y_train_h = humanDataset.readBatch(packetSize)
x_train_h, y_train_h = augmentations.appendTransforms(
x_train_h, y_train_h,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
logger.debug('reading human batch, memory used %f',
usedMemory())
x_train_nh, y_train_nh = nonHumanDataset.readBatch(
nonHumanPacketSize)
x_train_nh, y_train_nh = augmentations.appendTransforms(
x_train_nh, y_train_nh,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
logger.debug('reading nonHuman batch, memory used %f',
usedMemory())
x_train, y_train = HumanDatasetSequence.shuffleHumanNonHuman(
x_train_h, x_train_nh, y_train_h, y_train_nh)
activeLearningImagesNeedToRead += activeLearningImagesPerPacket
if activeLearningImagesNeedToRead - activeLearningImagesRead >= 1:
activeLearningImagesPacketSize = int(
activeLearningImagesNeedToRead -
activeLearningImagesRead)
remaining = len(
activeLearningDataset) - activeLearningDataset.index
print(
"activeLearningImagesPacketSize {}, index {}, remaining {}"
.format(activeLearningImagesPacketSize,
activeLearningDataset.index, remaining))
if activeLearningImagesPacketSize > remaining:
x_train_a1, y_train_a1 = activeLearningDataset.readBatch(
remaining)
x_train_a1, y_train_a1 = setFixedSize(
x_train_a1, y_train_a1)
x_train_a1, y_train_a1 = augmentations.appendTransforms(
x_train_a1, y_train_a1,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
activeLearningDataset.reset()
x_train_a2, y_train_a2 = activeLearningDataset.readBatch(
activeLearningImagesPacketSize - remaining)
x_train_a2, y_train_a2 = setFixedSize(
x_train_a2, y_train_a2)
x_train_a2, y_train_a2 = augmentations.appendTransforms(
x_train_a2, y_train_a2,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
x_train_a = x_train_a1 + x_train_a2
y_train_a = y_train_a1 + y_train_a2
else:
x_train_a, y_train_a = activeLearningDataset.readBatch(
activeLearningImagesPacketSize)
x_train_a, y_train_a = setFixedSize(
x_train_a, y_train_a)
x_train_a, y_train_a = augmentations.appendTransforms(
x_train_a, y_train_a,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
cv2.imshow("imageA", x_train_a[0])
cv2.imshow("maskA", y_train_a[0] * 255)
cv2.waitKey(1)
activeLearningImagesRead += activeLearningImagesPacketSize
print("activeLearningImagesRead {}".format(
activeLearningImagesRead))
x_train = np.concatenate((x_train, x_train_a))
y_train = np.concatenate((y_train, y_train_a))
del x_train_a
del y_train_a
else:
x_train = np.concatenate((x_train, ))
y_train = np.concatenate((y_train, ))
del x_train_h
del x_train_nh
del y_train_h
del y_train_nh
logger.debug('concatenate batches, memory used %f',
usedMemory())
x_train = preprocess_input(x_train)
# x_train = x_train / 255
logger.debug('preprocess x_train, memory used %f',
usedMemory())
saveModel = ((humanDataset.index + nonHumanDataset.index) %
50000) < (packetSize + nonHumanPacketSize)
logger.debug('start train on %d samples, memory used %f',
len(x_train), usedMemory())
model.fit(
x=x_train,
y=y_train,
batch_size=batchSize,
epochs=epoch + 1,
initial_epoch=epoch,
validation_data=(x_val, y_val),
)
if saveModel:
save_model(model, trainingDir, modelEncoder, packetIndex)
del x_train
del y_train
logger.debug('trained on %d samples, memory used %f',
humanDataset.index + nonHumanDataset.index,
usedMemory())
gc.collect()
# objgraph.show_most_common_types(limit=50)
# obj = objgraph.by_type('list')[1000]
# objgraph.show_backrefs(obj, max_depth=10)
x_train_h, y_train_h = humanDataset.readBatch(packetSize)
x_train_h, y_train_h = augmentations.appendTransforms(
x_train_h, y_train_h,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
x_train_nh, y_train_nh = nonHumanDataset.readBatch(
nonHumanPacketSize)
x_train_nh, y_train_nh = augmentations.appendTransforms(
x_train_nh, y_train_nh,
augmentations.train_transforms_after_resize,
augmentations.resize_transforms(imageSize))
x_train, y_train = HumanDatasetSequence.shuffleHumanNonHuman(
x_train_h, x_train_nh, y_train_h, y_train_nh)
x_train = np.concatenate((x_train, ))
y_train = np.concatenate((y_train, ))
del x_train_h
del x_train_nh
del y_train_h
del y_train_nh
x_train = preprocess_input(x_train)
# x_train = x_train / 255
model.fit(x=x_train,
y=y_train,
batch_size=batchSize,
epochs=1,
validation_data=(x_val, y_val))
save_model(model, trainingDir, modelEncoder, packets - 1)
del x_train
del y_train
logger.info('epoch %d is trained', epoch)
except Exception as e:
logger.error('Exception %s', str(e))
traceback.print_exc()
return
now = datetime.now()
dt_string = now.strftime("%Y/%m/%d %H:%M:%S")
modelPath = os.path.join(
trainingDir,
'u-net-{}_epoch{}_{}.tfmodel'.format(modelEncoder, epoch,
dt_string))
model.save(modelPath)
logger.info('model saved')