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)