def main():
parser = argparse.ArgumentParser(
description='GAN-based unsupervised segmentation train')
parser.add_argument('--unet_weights', type=str, default="")
parser.add_argument('--seed', type=int, default=2)
parser.add_argument('--val_images_dirs',
nargs='*',
type=str,
default=[None])
parser.add_argument('--val_masks_dirs',
nargs='*',
type=str,
default=[None])
args = parser.parse_args()
model = UNet().train().cuda()
model.load_state_dict(torch.load(args.unet_weights))
evaluate_all_wrappers(model, args.val_images_dirs, args.val_masks_dirs)
type=str,
default="output",
help="directory saving prediction results")
opt = parser.parse_args()
logfile = 'logs/predict/' + opt.dataset_name + '_' + opt.device + '.log'
sys.stdout = Logger(logfile)
print(opt)
device = torch.device(opt.device)
output_dir_unet = opt.output_dir + '/unet_segmentation/' + opt.dataset_name
os.makedirs(output_dir_unet, exist_ok=True)
unet_path = 'UNet/checkpoints/' + opt.unet_ckpt
model_unet = UNet(n_channels=1, n_classes=1).to(device=device)
model_unet.load_state_dict(torch.load(unet_path)['state_dict'])
model_unet.eval()
output_dir_yolo = opt.output_dir + '/yolo_detection/' + opt.dataset_name
os.makedirs(output_dir_yolo, exist_ok=True)
classes = ['vein']
yolo_path = 'YOLOv3/checkpoints/' + opt.yolo_ckpt
model_yolo = Darknet('YOLOv3/config/yolov3-custom.cfg',
img_size=416).to(device=device)
model_yolo.load_state_dict(torch.load(yolo_path))
model_yolo.eval()
image_folder = 'DATA/' + opt.dataset_name + '/imgs'
image_files = [x for x in os.listdir(image_folder)
if x.endswith('.jpg')] # only jpg files
for i, fn in enumerate(image_files):