提交 e5cabbb1 编写于 作者: C chenguowei01

update some

上级 30089840
...@@ -91,7 +91,7 @@ class OpticDiscSeg(Dataset): ...@@ -91,7 +91,7 @@ class OpticDiscSeg(Dataset):
elif self.mode == 'eval': elif self.mode == 'eval':
return im, label return im, label
if self.mode == 'test': if self.mode == 'test':
return im, im_info return im, im_info, image_path
def __len__(self): def __len__(self):
return len(self.file_list) return len(self.file_list)
...@@ -18,9 +18,11 @@ import os ...@@ -18,9 +18,11 @@ import os
from paddle.fluid.dygraph.base import to_variable from paddle.fluid.dygraph.base import to_variable
import numpy as np import numpy as np
import paddle.fluid as fluid import paddle.fluid as fluid
from paddle.fluid.dygraph.parallel import ParallelEnv
import cv2 import cv2
import tqdm import tqdm
from datasets import OpticDiscSeg
import transforms as T import transforms as T
import models import models
import utils import utils
...@@ -32,50 +34,58 @@ def parse_args(): ...@@ -32,50 +34,58 @@ def parse_args():
parser = argparse.ArgumentParser(description='Model training') parser = argparse.ArgumentParser(description='Model training')
# params of model # params of model
parser.add_argument('--model_name', parser.add_argument(
dest='model_name', '--model_name',
help="Model type for traing, which is one of ('UNet')", dest='model_name',
type=str, help="Model type for traing, which is one of ('UNet')",
default='UNet') type=str,
default='UNet')
# params of dataset # params of dataset
parser.add_argument('--data_dir', parser.add_argument(
dest='data_dir', '--data_dir',
help='The root directory of dataset', dest='data_dir',
type=str) help='The root directory of dataset',
parser.add_argument('--test_list', type=str)
dest='test_list', parser.add_argument(
help='Val list file of dataset', '--test_list',
type=str, dest='test_list',
default=None) help='Val list file of dataset',
parser.add_argument('--num_classes', type=str,
dest='num_classes', default=None)
help='Number of classes', parser.add_argument(
type=int, '--num_classes',
default=2) dest='num_classes',
help='Number of classes',
type=int,
default=2)
# params of prediction # params of prediction
parser.add_argument("--input_size", parser.add_argument(
dest="input_size", "--input_size",
help="The image size for net inputs.", dest="input_size",
nargs=2, help="The image size for net inputs.",
default=[512, 512], nargs=2,
type=int) default=[512, 512],
parser.add_argument('--batch_size', type=int)
dest='batch_size', parser.add_argument(
help='Mini batch size', '--batch_size',
type=int, dest='batch_size',
default=2) help='Mini batch size',
parser.add_argument('--model_dir', type=int,
dest='model_dir', default=2)
help='The path of model for evaluation', parser.add_argument(
type=str, '--model_dir',
default=None) dest='model_dir',
parser.add_argument('--save_dir', help='The path of model for evaluation',
dest='save_dir', type=str,
help='The directory for saving the inference results', default=None)
type=str, parser.add_argument(
default='./output/result') '--save_dir',
dest='save_dir',
help='The directory for saving the inference results',
type=str,
default='./output/result')
return parser.parse_args() return parser.parse_args()
...@@ -86,26 +96,19 @@ def mkdir(path): ...@@ -86,26 +96,19 @@ def mkdir(path):
os.makedirs(sub_dir) os.makedirs(sub_dir)
def infer(model, data_dir=None, test_list=None, model_dir=None, def infer(model, test_dataset=None, model_dir=None, save_dir='output'):
transforms=None):
ckpt_path = os.path.join(model_dir, 'model') ckpt_path = os.path.join(model_dir, 'model')
para_state_dict, opti_state_dict = fluid.load_dygraph(ckpt_path) para_state_dict, opti_state_dict = fluid.load_dygraph(ckpt_path)
model.set_dict(para_state_dict) model.set_dict(para_state_dict)
model.eval() model.eval()
added_saved_dir = os.path.join(args.save_dir, 'added') added_saved_dir = os.path.join(save_dir, 'added')
pred_saved_dir = os.path.join(args.save_dir, 'prediction') pred_saved_dir = os.path.join(save_dir, 'prediction')
logging.info("Start to predict...") logging.info("Start to predict...")
with open(test_list, 'r') as f: for im, im_info, im_path in tqdm.tqdm(test_dataset):
files = f.readlines() im = im[np.newaxis, ...]
for file in tqdm.tqdm(files):
file = file.strip()
im_file = os.path.join(data_dir, file)
im, im_info, _ = transforms(im_file)
im = np.expand_dims(im, axis=0)
im = to_variable(im) im = to_variable(im)
pred, _ = model(im, mode='test') pred, _ = model(im, mode='test')
pred = pred.numpy() pred = pred.numpy()
pred = np.squeeze(pred).astype('uint8') pred = np.squeeze(pred).astype('uint8')
...@@ -118,38 +121,41 @@ def infer(model, data_dir=None, test_list=None, model_dir=None, ...@@ -118,38 +121,41 @@ def infer(model, data_dir=None, test_list=None, model_dir=None,
h, w = im_info[k][0], im_info[k][1] h, w = im_info[k][0], im_info[k][1]
pred = pred[0:h, 0:w] pred = pred[0:h, 0:w]
im_file = im_path.replace(test_dataset.data_dir, '')
if im_file[0] == '/':
im_file = im_file[1:]
# save added image # save added image
added_image = utils.visualize(im_file, pred, weight=0.6) added_image = utils.visualize(im_path, pred, weight=0.6)
added_image_path = os.path.join(added_saved_dir, file) added_image_path = os.path.join(added_saved_dir, im_file)
mkdir(added_image_path) mkdir(added_image_path)
cv2.imwrite(added_image_path, added_image) cv2.imwrite(added_image_path, added_image)
# save prediction # save prediction
pred_im = utils.visualize(im_file, pred, weight=0.0) pred_im = utils.visualize(im_path, pred, weight=0.0)
pred_saved_path = os.path.join(pred_saved_dir, file) pred_saved_path = os.path.join(pred_saved_dir, im_file)
mkdir(pred_saved_path) mkdir(pred_saved_path)
cv2.imwrite(pred_saved_path, pred_im) cv2.imwrite(pred_saved_path, pred_im)
def main(args): def main(args):
env_info = get_environ_info()
places = fluid.CUDAPlace(ParallelEnv().dev_id) \
if env_info['place'] == 'cuda' and fluid.is_compiled_with_cuda() \
else fluid.CPUPlace()
with fluid.dygraph.guard(places): with fluid.dygraph.guard(places):
test_transforms = T.Compose([T.Resize(args.input_size), T.Normalize()]) test_transforms = T.Compose([T.Resize(args.input_size), T.Normalize()])
test_dataset = OpticDiscSeg(transforms=test_transforms, mode='test')
if args.model_name == 'UNet': if args.model_name == 'UNet':
model = models.UNet(num_classes=args.num_classes) model = models.UNet(num_classes=args.num_classes)
infer(model, infer(
data_dir=args.data_dir, model,
test_list=args.test_list, model_dir=args.model_dir,
model_dir=args.model_dir, test_dataset=test_dataset,
transforms=test_transforms) save_dir=args.save_dir)
if __name__ == '__main__': if __name__ == '__main__':
args = parse_args() args = parse_args()
env_info = get_environ_info()
if env_info['place'] == 'cpu':
places = fluid.CPUPlace()
else:
places = fluid.CUDAPlace(0)
main(args) main(args)
...@@ -116,6 +116,11 @@ def parse_args(): ...@@ -116,6 +116,11 @@ def parse_args():
help='Num workers for data loader', help='Num workers for data loader',
type=int, type=int,
default=0) default=0)
parser.add_argument(
'--do_eval',
dest='do_eval',
help='Eval while training',
action='store_true')
return parser.parse_args() return parser.parse_args()
...@@ -183,7 +188,7 @@ def train(model, ...@@ -183,7 +188,7 @@ def train(model,
"epoch_{}".format(epoch + 1)) "epoch_{}".format(epoch + 1))
if not os.path.isdir(current_save_dir): if not os.path.isdir(current_save_dir):
os.makedirs(current_save_dir) os.makedirs(current_save_dir)
fluid.save_dygraph(model_parallel.state_dict(), fluid.save_dygraph(model.state_dict(),
os.path.join(current_save_dir, 'model')) os.path.join(current_save_dir, 'model'))
if eval_dataset is not None: if eval_dataset is not None:
...@@ -215,12 +220,11 @@ def main(args): ...@@ -215,12 +220,11 @@ def main(args):
train_dataset = OpticDiscSeg(transforms=train_transforms, mode='train') train_dataset = OpticDiscSeg(transforms=train_transforms, mode='train')
eval_dataset = None eval_dataset = None
if args.val_list is not None: if args.do_eval:
eval_transforms = T.Compose( eval_transforms = T.Compose(
[T.Resize(args.input_size), [T.Resize(args.input_size),
T.Normalize()]) T.Normalize()])
eval_dataset = OpticDiscSeg( eval_dataset = OpticDiscSeg(transforms=eval_transforms, mode='eval')
transforms=train_transforms, mode='eval')
if args.model_name == 'UNet': if args.model_name == 'UNet':
model = models.UNet(num_classes=args.num_classes, ignore_index=255) model = models.UNet(num_classes=args.num_classes, ignore_index=255)
......
...@@ -19,25 +19,26 @@ import math ...@@ -19,25 +19,26 @@ import math
from paddle.fluid.dygraph.base import to_variable from paddle.fluid.dygraph.base import to_variable
import numpy as np import numpy as np
import paddle.fluid as fluid import paddle.fluid as fluid
from paddle.fluid.dygraph.parallel import ParallelEnv
from paddle.fluid.io import DataLoader from paddle.fluid.io import DataLoader
from paddle.fluid.dataloader import BatchSampler
from datasets import Dataset from datasets import OpticDiscSeg
import transforms as T import transforms as T
import models import models
import utils.logging as logging import utils.logging as logging
from utils import get_environ_info from utils import get_environ_info
from utils import ConfusionMatrix from utils import ConfusionMatrix
from utils import DistributedBatchSampler
def parse_args(): def parse_args():
parser = argparse.ArgumentParser(description='Model training') parser = argparse.ArgumentParser(description='Model evaluation')
# params of model # params of model
parser.add_argument( parser.add_argument(
'--model_name', '--model_name',
dest='model_name', dest='model_name',
help="Model type for traing, which is one of ('UNet')", help="Model type for evaluation, which is one of ('UNet')",
type=str, type=str,
default='UNet') default='UNet')
...@@ -97,28 +98,28 @@ def evaluate(model, ...@@ -97,28 +98,28 @@ def evaluate(model,
model.set_dict(para_state_dict) model.set_dict(para_state_dict)
model.eval() model.eval()
batch_sampler = DistributedBatchSampler( batch_sampler = BatchSampler(
eval_dataset, batch_size=batch_size, shuffle=True, drop_last=False) eval_dataset, batch_size=batch_size, shuffle=False, drop_last=False)
loader = DataLoader( loader = DataLoader(
eval_dataset, eval_dataset,
batch_sampler=batch_sampler, batch_sampler=batch_sampler,
places=places, places=places,
return_list=True, return_list=True,
) )
total_steps = math.ceil(eval_dataset.num_samples * 1.0 / batch_size) total_steps = math.ceil(len(eval_dataset) * 1.0 / batch_size)
conf_mat = ConfusionMatrix(num_classes, streaming=True) conf_mat = ConfusionMatrix(num_classes, streaming=True)
logging.info( logging.info(
"Start to evaluating(total_samples={}, total_steps={})...".format( "Start to evaluating(total_samples={}, total_steps={})...".format(
eval_dataset.num_samples, total_steps)) len(eval_dataset), total_steps))
for step, data in enumerate(loader): for step, data in enumerate(loader):
images = data[0] images = data[0]
labels = data[1].astype('int64') labels = data[1].astype('int64')
pred, _ = model(images, labels, mode='eval') pred, _ = model(images, mode='eval')
pred = pred.numpy() pred = pred.numpy()
labels = labels.numpy()
mask = labels != ignore_index mask = labels != ignore_index
conf_mat.calculate(pred=pred, label=labels, ignore=mask) conf_mat.calculate(pred=pred, label=labels, ignore=mask)
_, iou = conf_mat.mean_iou() _, iou = conf_mat.mean_iou()
...@@ -128,7 +129,7 @@ def evaluate(model, ...@@ -128,7 +129,7 @@ def evaluate(model,
category_iou, miou = conf_mat.mean_iou() category_iou, miou = conf_mat.mean_iou()
category_acc, macc = conf_mat.accuracy() category_acc, macc = conf_mat.accuracy()
logging.info("[EVAL] #image={} acc={:.4f} IoU={:.4f}".format( logging.info("[EVAL] #image={} acc={:.4f} IoU={:.4f}".format(
eval_dataset.num_samples, macc, miou)) len(eval_dataset), macc, miou))
logging.info("[EVAL] Category IoU: " + str(category_iou)) logging.info("[EVAL] Category IoU: " + str(category_iou))
logging.info("[EVAL] Category Acc: " + str(category_acc)) logging.info("[EVAL] Category Acc: " + str(category_acc))
logging.info("[EVAL] Kappa:{:.4f} ".format(conf_mat.kappa())) logging.info("[EVAL] Kappa:{:.4f} ".format(conf_mat.kappa()))
...@@ -136,20 +137,12 @@ def evaluate(model, ...@@ -136,20 +137,12 @@ def evaluate(model,
def main(args): def main(args):
env_info = get_environ_info() env_info = get_environ_info()
if env_info['place'] == 'cpu': places = fluid.CUDAPlace(ParallelEnv().dev_id) \
places = fluid.CPUPlace() if env_info['place'] == 'cuda' and fluid.is_compiled_with_cuda() \
else: else fluid.CPUPlace()
places = fluid.CUDAPlace(0)
with fluid.dygraph.guard(places): with fluid.dygraph.guard(places):
eval_transforms = T.Compose([T.Resize(args.input_size), T.Normalize()]) eval_transforms = T.Compose([T.Resize(args.input_size), T.Normalize()])
eval_dataset = Dataset( eval_dataset = OpticDiscSeg(transforms=eval_transforms, mode='eval')
data_dir=args.data_dir,
file_list=args.val_list,
transforms=eval_transforms,
num_workers='auto',
buffer_size=100,
parallel_method='thread',
shuffle=False)
if args.model_name == 'UNet': if args.model_name == 'UNet':
model = models.UNet(num_classes=args.num_classes) model = models.UNet(num_classes=args.num_classes)
...@@ -157,6 +150,7 @@ def main(args): ...@@ -157,6 +150,7 @@ def main(args):
evaluate( evaluate(
model, model,
eval_dataset, eval_dataset,
places=places,
model_dir=args.model_dir, model_dir=args.model_dir,
num_classes=args.num_classes, num_classes=args.num_classes,
batch_size=args.batch_size) batch_size=args.batch_size)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册