def main():
args = build_parser().parse_args()
image_size = [args.img_height, args.img_width]
# config = tf.ConfigProto()
# config.gpu_options.per_process_gpu_memory_fraction = 1.0
# sess = tf.Session(config=config)
sess = tf.Session()
unet = Unet(input_shape=image_size,
sess=sess,
filter_num=args.filter_num,
batch_norm=args.batch_norm)
unet.build_net()
if args.checkpoint_path:
unet.load_weights(args.checkpoint_path)
images, masks = read_data(args.train_dir,
args.train_mask_dir,
n_images=args.n_images,
image_size=image_size)
val_images, val_masks = read_data(args.val_dir,
args.val_mask_dir,
n_images=args.n_images // 4,
image_size=image_size)
unet.train(images=images,
masks=masks,
val_images=val_images,
val_masks=val_masks,
epochs=args.epochs,
batch_size=args.batch_size,
learning_rate=args.learning_rate,
dice_loss=args.dice_loss,
always_save=args.always_save)
def get_value():
modelpath = A.get()
picinit = B.get()
outpath = F.get()
#pic = list(pic)
pic = picinit.split(' ')
sign1 = sign.get()
num_class1 = num_class.get()
if not modelpath or not pic:
change.set('请选择文件')
if modelpath and pic:
change.set('执行中...')
root.update()
#allpic = glob.glob(os.path.join(pic,'*.tif'))
lastpic = pic[-1]
model = Unet((256, 256, 3), num_class1)
#model = myunet((256,256,3),num_class1)
model.load_weights(modelpath)
d, n = os.path.split(lastpic)
lastpic_save = os.path.join(outpath + '/' + n) #最后一个文件
# delete_path = os.path.join(d,'result')#保存文件目录
# if os.path.exists(delete_path):
# pp = os.listdir(delete_path)
# for x in pp:
# delete = os.path.join(delete_path,x)
# os.remove(delete)
# os.removedirs(delete_path)
W = P(num_class1)
W.main_p(model, pic, outpath, changes=sign1)
if os.path.exists(lastpic_save):
change.set('识别完成！')
os.startfile(outpath)
def main():
args = build_parser().parse_args()
assert args.checkpoint_path
result_dir = args.result_dir
checkpoint_path = args.checkpoint_path
test_dir = args.test_dir
n_imgs = args.n_images
image_size = [args.img_height, args.img_width]
sess = tf.Session()
unet = Unet(input_shape=image_size,
sess=sess,
filter_num=args.filter_num,
batch_norm=args.batch_norm)
unet.build_net(is_train=False)
unet.load_weights(checkpoint_path)
img_names = os.listdir(test_dir)
img_names.sort()
mask_names = None
total_dice = None
if args.mask_dir:
mask_names = os.listdir(args.mask_dir)
mask_names.sort()
total_dice = 0
if n_imgs <= 0:
n_imgs = len(img_names)
for i in range(n_imgs):
print('%s %d/%d' % (img_names[i], i, n_imgs))
img_mat = read_car_img(os.path.join(test_dir, img_names[i]),
image_size=image_size)
img_mat = np.expand_dims(img_mat, axis=0)
if mask_names:
mask_mat = read_mask_img(os.path.join(args.mask_dir,
mask_names[i]),
image_size=image_size)
mask_mat = np.expand_dims(mask_mat, axis=0)
res, dice = unet.predict_test(img_mat, mask_mat)
dice = np.mean(dice)
print('Dice coefficient:%.6f' % dice)
total_dice += dice
else:
res = unet.predict(img_mat)
if args.result_dir:
res = res.reshape(image_size)
misc.imsave(os.path.join(result_dir, img_names[i]), res)
if total_dice:
print('Average Dice coefficient:%.6f' % (total_dice / n_imgs))
w = epoch // 10
lr = init_lr / (lr_epoch_decay**w)
if lr < 1e-10:
lr = 1e-10
return lr
callback = LearningRateScheduler(get_lr)
callbacks.append(callback)
return callbacks
model = Unet(input_shape, num)
model.summary()
if os.path.exists(mode1_path): #继续训练
model.load_weights(mode1_path)
callbacks = make_callbacks()
path = r'F:\cmm\yumi\pic'
train_set, val_set = get_train_val(path, 'tif')
train = train_data(train_set, bitchs)
val = val_data(val_set, bitchs)
alltrain = len(train_set)
allval = len(val_set)
loss_f = loss.mean_iou
metrics = [loss_f]
model.compile(optimizer='adam',
loss='categorical_crossentropy',
metrics=metrics)
A = model.fit_generator(generator=train,
steps_per_epoch=alltrain // bitchs,
epochs=epochs,
def predict(model_name,
mask_save=False,
combine_save=True,
crf_treat=True,
keep_path_struct=True):
'''
参数说明：
mask_save：设定是否单独保存预测结果
combine_save：设定是否保存与原图的融合结果
crf_treat：设定是否将预测结果进行CRF后处理---该处理目的是使得预测边缘平滑化，减少误判
keep_path_struct：设定预测结果的保存，是否遵循原预测文件的架构。False代表统一存放到设定的文件夹下
'''
def real_pred(imgs):
for nums, jpg in enumerate(imgs):
img_name = os.path.basename(jpg)
img_dir = os.path.dirname(jpg)
if keep_path_struct:
save_path = img_dir.replace(test_path, result_path) + "/"
if not os.path.exists(save_path):
os.makedirs(save_path)
else:
save_path = result_path + "/"
img = cv2.imread(jpg)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #通道转换
old_img = copy.deepcopy(img)
orininal_h, orininal_w = old_img.shape[:2]
img = letterbox_image(img, (HEIGHT, WIDTH)) #先处理成原始模型训练的尺寸大小
img = img / 255
img = img.reshape(-1, HEIGHT, WIDTH, 3)
pr = model.predict(img)[0]
if crf_treat:
pr = pr.reshape((-1, HEIGHT, WIDTH, NCLASSES))
temp_img = letterbox_image(
old_img, (HEIGHT, WIDTH)) #原图需要先转成训练时指定的图片大小
temp_img = temp_img[np.newaxis, ...] #需增加一个维度，方便输入CRF处理
pr = dense_crf(pr, img=temp_img,
n_classes=NCLASSES) #增加CRF后处理
pr = pr.reshape(
(HEIGHT, WIDTH,
NCLASSES)).argmax(axis=-1) #将刚检测出来的图片，转成单通道的标签图片大小的数据格式
pr = np.uint8(pr)
pr = cv2.resize(pr, (orininal_w, orininal_h))
for num in range(1, NCLASSES): #逐个将每个类别的像素值，计入各自的列表中
num_cal = Counter(
pr[pr == num])[num] #Counter的结果类似一个字典，直接取其键值
areas[classes[num]].append(num_cal)
if mask_save:
mask = label2rgb(pr,
n_labels=len(classes),
colormap=colors,
mask_save=mask_save) #给mask上色, 注释掉则变单通道图
B, G, R = cv2.split(mask)
A = np.zeros(B.shape, dtype=B.dtype) #增加alpha通道，方便前端调用
A[B > 0] = 255
mask = cv2.merge((B, G, R, A))
cv2.imwrite(save_path + img_name[:-4] + ".png", mask)
if combine_save:
image = draw_label(pr, old_img, classes, colormap=colors)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
cv2.imwrite(save_path + img_name, image)
print('文件夹%s的第%d张图片检测完毕' % (img_dir, nums + 1))
if model_name == "deeplabv3":
model = Deeplabv3(input_shape=(HEIGHT, WIDTH, 3),
classes=NCLASSES,
bone_name=bone_name)
elif model_name == 'unet':
model = Unet(input_shape=(HEIGHT, WIDTH, 3),
classes=NCLASSES,
bone_name=bone_name)
elif model_name == 'pspnet':
model = PSPnet(input_shape=(HEIGHT, WIDTH, 3),
classes=NCLASSES,
bone_name=bone_name)
log_name = os.path.basename(find_last(log_path))
print('-' * 100)
print('Loading log_name of %s' % log_name)
print('-' * 100)
model.load_weights(find_last(log_path), by_name=True)
#需要根据类别数，来随机设定不同类别的mask对应的颜色
colors = label_colormap(NCLASSES) #定义函数来确定不同类别的颜色，相对颜色识别度较高。
total_area = 0 #初始化总面积
areas = {} #设置空字典，用于存放每张图片各个类别的面积，并最终求和
for num in range(1, NCLASSES): #先初始化每个类别，用于存放每张图得到的各类别像素点个数
areas[classes[num]] = []
print('开始检测...')
for roots, _, _ in os.walk(test_path): #递归遍历多级子目录
imgs = glob.glob(roots + "/*.jpg") #这里需注意，如果原图片是png格式的，需更改成.png
real_pred(imgs)
#统计各类别的面积
f = open(result_path + '/area.txt', 'w')
for num in range(1, NCLASSES):
area = sum(areas[classes[num]])
total_area += area
f.write('%s area is : %.2f m2 \n' %
(classes[num], area * area_per_pixel))
f.write('total area is : %.2f m2' % (total_area * area_per_pixel))
f.close()
t_end = time.time()
time_total = pd.to_timedelta(t_end - t_start, unit='s')
print('%s_%s模型检测结束, 检测图片%s张, 结束时间为%s, 总耗时%s秒' %
(bone_name, model_name, len(areas[classes[num]]),
time.strftime('%Y-%m-%d %H:%M:%S'), time_total.round('s')))
pic = stretch(pic)
pic = pic.astype(np.float32)
#pic = img_to_array(pic)
y_probs = make_prediction_img(
pic, 256, 8,
lambda xx: predict_x(xx, model)) # 数据，目标大小，批次大小，返回每次识别的
y_preds = np.argmax(y_probs, axis=2)
d, n = os.path.split(one_path)
t0 = time.time()
change = y_preds.astype(np.uint8)
outpath = os.path.join(d, 'result')
if not os.path.exists(outpath):
os.makedirs(outpath)
CreatTf(one_path, change, outpath) # 添加坐标系
img_out = np.zeros(change.shape + (3, ))
for i in range(num_class):
img_out[change == i, :] = COLOR_DICT[i] #对应上色
change = img_out / 255
save_file = os.path.join(outpath, n[:-4] + '_color' + '.png')
skimage.io.imsave(save_file, change)
print('预测耗费时间: %0.2f(min).' % ((time.time() - t0) / 60))
if __name__ == '__main__':
model = Unet((256, 256, 3), num_class)
p = r'E:\buildingone\output\YMDD.h5' # 说明权重所在位置
print("网络参数来自: '%s'." % p)
model.load_weights(p)
path = r'E:\mynet\end'
main_p(model, path, changes=False)
def train_model(train_data,
val_data,
num_train,
num_val,
epochs,
callback=True): #callback用于是否记录训练。
if model_name == "deeplabv3":
model = Deeplabv3((HEIGHT, WIDTH, 3), NCLASSES, bone_name)
elif model_name == 'unet':
model = Unet((HEIGHT, WIDTH, 3), NCLASSES, bone_name)
elif model_name == 'pspnet':
model = PSPnet((HEIGHT, WIDTH, 3), NCLASSES, bone_name)
epoch = 0
if init_with == 'first': #init_with selection = ['first', 'last'] 选择模型的训练模式
print('-' * 100)
# model.load_weights(weights_path, by_name=True, skip_mismatch=True) #加载预训练模型
print('开始从头训练模型...')
else:
model.load_weights(find_last(log_path), by_name=True)
epoch = int(
os.path.basename(find_last(log_path))[:-3].split('_')[-1])
epochs = epoch + epochs #这样可以确保重新训练时，设置的epochs为实际训练的轮数。
print('-' * 100)
print('成功加载最新模型, 重新从第%s轮开始训练...' % epoch)
#保存训练过程
tbCallBack = TensorBoard(log_dir=log_path + "records/",
histogram_freq=0,
write_graph=True,
write_images=True)
# 保存的方式，1次epoch保存一次
checkpoint_path = os.path.join(
log_path, '%s_%s_*epoch*.h5' % (model_name, bone_name))
checkpoint_path = checkpoint_path.replace("*epoch*", "{epoch:04d}")
checkpoint_period = ModelCheckpoint(checkpoint_path,
monitor='val_loss',
save_weights_only=True,
save_best_only=True,
period=1)
# 学习率下降的方式，val_loss五次不下降就下降学习率继续训练
reduce_lr = ReduceLROnPlateau(monitor='val_loss',
factor=0.5,
patience=5,
verbose=1)
# 是否需要早停，当val_loss一直不下降的时候意味着模型基本训练完毕，可以停止
early_stopping = EarlyStopping(monitor='val_loss',
min_delta=0,
patience=15,
verbose=1)
callbacks = [checkpoint_period, reduce_lr, early_stopping]
if callback:
tbCallBack.set_model(model)
callbacks.append(tbCallBack)
print('训练的样本量为：{} 张图片, 验证集的样本量为：{} 张图片，'.format(num_train, num_val))
print('每份数据集大小为：{} 张图片, 图片大小为: {}'.format(batch_size, (HEIGHT, WIDTH)))
print('以%s为主干的%s模型正式开始训练，请耐心等待，并注意提示...' %
(bone_name, model_name)) #开始训练
print('-' * 100)
model.compile(loss=focal_loss,
optimizer=Adam(lr=1e-3),
metrics=['accuracy'])
model.fit_generator(generate_arrays_from_file(train_data, batch_size),
steps_per_epoch=max(1, num_train // batch_size),
validation_data=generate_arrays_from_file(
val_data, batch_size),
validation_steps=max(1, num_val // batch_size),
epochs=epochs,
initial_epoch=epoch,
shuffle=True,
callbacks=callbacks)
