def output_point_cloud_ply(xyzs, name, output_folder): if not os.path.exists(output_folder): mkdir_p(output_folder) print('write: ' + os.path.join(output_folder, name + '.ply')) with open(os.path.join(output_folder, name + '.ply'), 'w') as f: pn = xyzs.shape[0] f.write('ply\n') f.write('format ascii 1.0\n') f.write('element vertex %d\n' % (pn)) f.write('property float x\n') f.write('property float y\n') f.write('property float z\n') f.write('end_header\n') for i in range(pn): f.write('%f %f %f\n' % (xyzs[i][0], xyzs[i][1], xyzs[i][2]))
def test(test_loader, model, args, save_result=False, best_epoch=None): global device model.eval() # switch to test mode loss_meter = AverageMeter() outdir = args.checkpoint.split('/')[1] for data in test_loader: data = data.to(device) with torch.no_grad(): if args.arch == 'masknet': mask_pred = model(data) mask_gt = data.mask.unsqueeze(1) loss = torch.nn.functional.binary_cross_entropy_with_logits( mask_pred, mask_gt.float(), reduction='mean') elif args.arch == 'jointnet': data_displacement = model(data) y_pred = data_displacement + data.pos loss = 0.0 for i in range(len(torch.unique(data.batch))): y_gt_sample = data.y[data.batch == i, :] y_gt_sample = y_gt_sample[:data.num_joint[i], :] y_pred_sample = y_pred[data.batch == i, :] loss += chamfer_distance_with_average( y_pred_sample.unsqueeze(0), y_gt_sample.unsqueeze(0)) loss_meter.update(loss.item()) if save_result: output_folder = 'results/{:s}/best_{:d}/'.format( outdir, best_epoch) if not os.path.exists(output_folder): mkdir_p(output_folder) if args.arch == 'masknet': mask_pred = torch.sigmoid(mask_pred) for i in range(len(torch.unique(data.batch))): mask_pred_sample = mask_pred[data.batch == i] np.save( os.path.join( output_folder, str(data.name[i].item()) + '_attn.npy'), mask_pred_sample.data.cpu().numpy()) else: for i in range(len(torch.unique(data.batch))): y_pred_sample = y_pred[data.batch == i, :] output_point_cloud_ply( y_pred_sample, name=str(data.name[i].item()), output_folder='results/{:s}/best_{:d}/'.format( outdir, best_epoch)) return loss_meter.avg
def test(test_loader, model, args, save_result=False): global device model.eval() # switch to test mode loss_meter = AverageMeter() outdir = args.checkpoint.split('/')[-1] for data in test_loader: data = data.to(device) with torch.no_grad(): skin_pred = model(data) skin_gt = data.skin_label[:, 0:args.nearest_bone] loss_mask_batch = data.loss_mask.float()[:, 0:args.nearest_bone] skin_gt = skin_gt * loss_mask_batch skin_gt = skin_gt / (torch.sum(torch.abs(skin_gt), dim=1, keepdim=True) + 1e-8) vert_mask = (torch.abs(skin_gt.sum(dim=1) - 1.0) < 1e-8).float() loss = cross_entropy_with_probs(skin_pred, skin_gt, reduction='none') loss = (loss * loss_mask_batch * vert_mask.unsqueeze(1)).sum() / (loss_mask_batch * vert_mask.unsqueeze(1)).sum() loss_meter.update(loss.item()) if save_result: output_folder = 'results/{:s}/'.format(outdir) if not os.path.exists(output_folder): mkdir_p(output_folder) for i in range(len(torch.unique(data.batch))): print('output result for model {:d}'.format(data.name[i].item())) skin_pred_i = skin_pred[data.batch == i] bone_names = get_bone_names(os.path.join(args.test_folder, "{:d}_skin.txt".format(data.name[i].item()))) tpl_e = np.loadtxt(os.path.join(args.test_folder, "{:d}_tpl_e.txt".format(data.name[i].item()))).T loss_mask_sample = data.loss_mask.float()[data.batch == i, 0:args.nearest_bone] skin_pred_i = torch.softmax(skin_pred_i, dim=1) skin_pred_i = skin_pred_i * loss_mask_sample skin_nn_i = data.skin_nn[data.batch == i, 0:args.nearest_bone] skin_pred_asarray = np.zeros((len(skin_pred_i), len(bone_names))) for v in range(len(skin_pred_i)): for nn_id in range(len(skin_nn_i[v, :])): skin_pred_asarray[v, skin_nn_i[v, nn_id]] = skin_pred_i[v, nn_id] skin_pred_asarray = post_filter(skin_pred_asarray, tpl_e, num_ring=1) skin_pred_asarray[skin_pred_asarray < np.max(skin_pred_asarray, axis=1, keepdims=True) * 0.5] = 0.0 skin_pred_asarray = skin_pred_asarray / (skin_pred_asarray.sum(axis=1, keepdims=True) + 1e-10) with open(os.path.join(output_folder, "{:d}_bone_names.txt".format(data.name[i].item())), 'w') as fout: for bone_name in bone_names: fout.write("{:s} {:s}\n".format(bone_name[0], bone_name[1])) np.save(os.path.join(output_folder, "{:d}_full_pred.npy".format(data.name[i].item())), skin_pred_asarray) skel_filename = os.path.join(args.info_folder, "{:d}.txt".format(data.name[i].item())) output_rigging(skel_filename, skin_pred_asarray, output_folder, data.name[i].item()) return loss_meter.avg
def test(test_loader, model, args, save_result=False, best_epoch=None): global device model.eval() # switch to test mode if save_result: output_folder = 'results/{:s}/best_{:d}/'.format( args.checkpoint.split('/')[1], best_epoch) if not os.path.exists(output_folder): mkdir_p(output_folder) loss_meter = AverageMeter() for data in test_loader: data = data.to(device) with torch.no_grad(): pre_label, label = model(data) loss = torch.nn.functional.binary_cross_entropy_with_logits( pre_label, label.float()) if save_result: connect_prob = torch.sigmoid(pre_label) accumulate_start_id = 0 for i in range(len(torch.unique(data.batch))): pair_idx = data.pairs[ accumulate_start_id:accumulate_start_id + data.num_pair[i]].long() connect_prob_i = connect_prob[ accumulate_start_id:accumulate_start_id + data.num_pair[i]] accumulate_start_id += data.num_pair[i] cost_matrix = np.zeros( (data.num_joint[i], data.num_joint[i])) pair_idx = pair_idx.data.cpu().numpy() cost_matrix[pair_idx[:, 0], pair_idx[:, 1]] = connect_prob_i.data.cpu().numpy( ).squeeze() cost_matrix = 1 - cost_matrix print('saving: {:s}'.format( str(data.name[i].item()) + '_cost.npy')) np.save( os.path.join(output_folder, str(data.name[i].item()) + '_cost.npy'), cost_matrix) loss_meter.update(loss.item(), n=len(torch.unique(data.batch))) return loss_meter.avg
def main(args): global device lowest_loss = 1e20 # create checkpoint dir and log dir if not isdir(args.checkpoint): print("Create new checkpoint folder " + args.checkpoint) mkdir_p(args.checkpoint) if not args.resume: if isdir(args.logdir): shutil.rmtree(args.logdir) mkdir_p(args.logdir) # create model model = PairCls() model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # optionally resume from a checkpoint if args.resume: if isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) args.start_epoch = checkpoint['epoch'] lowest_loss = checkpoint['lowest_loss'] print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) train_loader = DataLoader(GraphDataset(root=args.train_folder), batch_size=args.train_batch, shuffle=True) val_loader = DataLoader(GraphDataset(root=args.val_folder), batch_size=args.test_batch, shuffle=False) test_loader = DataLoader(GraphDataset(root=args.test_folder), batch_size=args.test_batch, shuffle=False) if args.evaluate: print('\nEvaluation only') test_loss = test(test_loader, model, args, save_result=True, best_epoch=args.start_epoch) print('test_loss {:8f}'.format(test_loss)) return lr = args.lr logger = SummaryWriter(log_dir=args.logdir) for epoch in range(args.start_epoch, args.epochs): print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr)) lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule, args.gamma) train_loss = train(train_loader, model, optimizer, args) val_loss = test(val_loader, model, args) test_loss = test(test_loader, model, args, best_epoch=epoch + 1) print('Epoch{:d}. train_loss: {:.6f}.'.format(epoch, train_loss)) print('Epoch{:d}. val_loss: {:.6f}.'.format(epoch, val_loss)) print('Epoch{:d}. test_loss: {:.6f}.'.format(epoch, test_loss)) # remember best acc and save checkpoint is_best = val_loss < lowest_loss lowest_loss = min(val_loss, lowest_loss) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'lowest_loss': lowest_loss, 'optimizer': optimizer.state_dict() }, is_best, checkpoint=args.checkpoint) info = { 'train_loss': train_loss, 'val_loss': val_loss, 'test_loss': test_loss } for tag, value in info.items(): logger.add_scalar(tag, value, epoch + 1) print("=> loading checkpoint '{}'".format( os.path.join(args.checkpoint, 'model_best.pth.tar'))) checkpoint = torch.load(os.path.join(args.checkpoint, 'model_best.pth.tar')) best_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint '{}' (epoch {})".format( os.path.join(args.checkpoint, 'model_best.pth.tar'), best_epoch)) test_loss = test(test_loader, model, args, save_result=True, best_epoch=best_epoch) print('Best epoch:\n test_loss {:8f}'.format(test_loss))
def one_process(dataset_folder, start_id, end_id): model_list = np.loadtxt(os.path.join(dataset_folder, 'model_list.txt'), dtype=int) model_list = model_list[start_id:end_id] remesh_obj_folder = os.path.join(dataset_folder, "obj_remesh") mkdir_p(os.path.join(dataset_folder, "volumetric_geodesic/")) for model_id in model_list: print(model_id) if os.path.exists( os.path.join( dataset_folder, "volumetric_geodesic/{:d}_volumetric_geo.npy".format( model_id))): continue remeshed_obj_filename = os.path.join( dataset_folder, 'obj_remesh/{:d}.obj'.format(model_id)) ori_obj_filename = os.path.join(dataset_folder, 'obj/{:d}.obj'.format(model_id)) info_filename = os.path.join(dataset_folder, 'rig_info/{:d}.txt'.format(model_id)) pts = np.array( o3d.io.read_triangle_mesh( os.path.join(remesh_obj_folder, '{:d}.obj'.format(model_id))).vertices) mesh_remesh = trimesh.load(remeshed_obj_filename) mesh_ori = trimesh.load(ori_obj_filename) rig_info = Info(info_filename) bones, bone_name, _ = get_bones(rig_info) origins, ends, pts_bone_dist = pts2line(pts, bones) if os.path.exists( os.path.join( dataset_folder, "volumetric_geodesic/{:d}_visibility_raw.npy".format( model_id))): pts_bone_visibility = np.load( os.path.join( dataset_folder, "volumetric_geodesic/{:d}_visibility_raw.npy".format( model_id))) else: # pick one mesh with fewer faces to speed up if len(mesh_remesh.faces) < len(mesh_ori.faces): trimesh.repair.fix_normals(mesh_remesh) pts_bone_visibility = calc_pts2bone_visible_mat( mesh_remesh, origins, ends) else: trimesh.repair.fix_normals(mesh_ori) pts_bone_visibility = calc_pts2bone_visible_mat( mesh_ori, origins, ends) pts_bone_visibility = pts_bone_visibility.reshape( len(bones), len(pts)).transpose() #np.save(os.path.join(dataset_folder, "volumetric_geodesic/{:d}_visibility_raw.npy".format(model_id)), pts_bone_visibility) pts_bone_dist = pts_bone_dist.reshape(len(bones), len(pts)).transpose() # remove visible points which are too far if os.path.exists( os.path.join( dataset_folder, "volumetric_geodesic/{:d}_visibility_filtered.npy".format( model_id))): pts_bone_visibility = np.load( os.path.join( dataset_folder, "volumetric_geodesic/{:d}_visibility_filtered.npy".format( model_id))) else: for b in range(pts_bone_visibility.shape[1]): visible_pts = np.argwhere( pts_bone_visibility[:, b] == 1).squeeze(1) if len(visible_pts) == 0: continue threshold_b = np.percentile(pts_bone_dist[visible_pts, b], 15) pts_bone_visibility[pts_bone_dist[:, b] > 1.3 * threshold_b, b] = False #np.save(os.path.join(dataset_folder, "volumetric_geodesic/{:d}_visibility_filtered.npy".format(model_id)), pts_bone_visibility) mesh = o3d.io.read_triangle_mesh( os.path.join(remesh_obj_folder, '{:d}.obj'.format(model_id))) surface_geodesic = calc_surface_geodesic(mesh) visible_matrix = np.zeros(pts_bone_visibility.shape) visible_matrix[np.where( pts_bone_visibility == 1)] = pts_bone_dist[np.where( pts_bone_visibility == 1)] euc_dist = np.sqrt( np.sum((pts[np.newaxis, ...] - pts[:, np.newaxis, :])**2, axis=2)) for c in range(visible_matrix.shape[1]): unvisible_pts = np.argwhere(pts_bone_visibility[:, c] == 0).squeeze(1) visible_pts = np.argwhere(pts_bone_visibility[:, c] == 1).squeeze(1) if len(visible_pts) == 0: visible_matrix[:, c] = pts_bone_dist[:, c] continue for r in unvisible_pts: dist1 = np.min(surface_geodesic[r, visible_pts]) nn_visible = visible_pts[np.argmin( surface_geodesic[r, visible_pts])] if np.isinf(dist1): visible_matrix[r, c] = 8.0 + pts_bone_dist[r, c] else: visible_matrix[r, c] = dist1 + visible_matrix[nn_visible, c] np.save( os.path.join( dataset_folder, "volumetric_geodesic/{:d}_volumetric_geo.npy".format( model_id)), visible_matrix)
def main(args): global device best_acc = 0.0 # create checkpoint dir and log dir if not isdir(args.checkpoint): print("Create new checkpoint folder " + args.checkpoint) mkdir_p(args.checkpoint) if not args.resume: if isdir(args.logdir): shutil.rmtree(args.logdir) mkdir_p(args.logdir) # create model model = ROOTNET() model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # optionally resume from a checkpoint if args.resume: if isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_acc = checkpoint['best_acc'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) lr = optimizer.param_groups[0]['lr'] print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) train_loader = DataLoader(GraphDataset(root=args.train_folder), batch_size=args.train_batch, shuffle=True, follow_batch=['joints']) val_loader = DataLoader(GraphDataset(root=args.val_folder), batch_size=args.test_batch, shuffle=False, follow_batch=['joints']) test_loader = DataLoader(GraphDataset(root=args.test_folder), batch_size=args.test_batch, shuffle=False, follow_batch=['joints']) if args.evaluate: print('\nEvaluation only') test_loss, test_acc = test(test_loader, model) print('test_loss {:.8f}. test_acc: {:.6f}'.format(test_loss, test_acc)) return scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, args.schedule, gamma=args.gamma) logger = SummaryWriter(log_dir=args.logdir) for epoch in range(args.start_epoch, args.epochs): lr = scheduler.get_last_lr() print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr[0])) train_loss = train(train_loader, model, optimizer, args) val_loss, val_acc = test(val_loader, model) test_loss, test_acc = test(test_loader, model) scheduler.step() print('Epoch{:d}. train_loss: {:.6f}.'.format(epoch + 1, train_loss)) print('Epoch{:d}. val_loss: {:.6f}. val_acc: {:.6f}'.format( epoch + 1, val_loss, val_acc)) print('Epoch{:d}. test_loss: {:.6f}. test_acc: {:.6f}'.format( epoch + 1, test_loss, test_acc)) # remember best acc and save checkpoint is_best = val_acc > best_acc best_acc = max(val_acc, best_acc) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'best_acc': best_acc, 'optimizer': optimizer.state_dict() }, is_best, checkpoint=args.checkpoint) info = { 'train_loss': train_loss, 'val_loss': val_loss, 'val_accuracy': val_acc, 'test_loss': test_loss, 'test_accuracy': test_acc } for tag, value in info.items(): logger.add_scalar(tag, value, epoch + 1) print("=> loading checkpoint '{}'".format( os.path.join(args.checkpoint, 'model_best.pth.tar'))) checkpoint = torch.load(os.path.join(args.checkpoint, 'model_best.pth.tar')) best_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint '{}' (epoch {})".format( os.path.join(args.checkpoint, 'model_best.pth.tar'), best_epoch)) test_loss, test_acc = test(test_loader, model) print('Best epoch:\n test_loss {:8f} test_acc {:8f}'.format( test_loss, test_acc))
############################################################################### ############################################################################### import os import numpy as np from utils.os_utils import mkdir_p from utils.job_utils import JobScheduler from felpy.model.wavefront import Wavefront indir = "/gpfs/exfel/data/user/guestt/labwork/dCache/NanoKB-Pulse/EHC/" outdir1 = "/gpfs/exfel/data/group/spb-sfx/user/guestt/h5/NanoKB-Pulse/data/EHC/integrated/" outdir2 = "/gpfs/exfel/data/group/spb-sfx/user/guestt/h5/NanoKB-Pulse/data/EHC/cmplx/" mkdir_p(outdir1) mkdir_p(outdir2) def launch(): """ This part launches the jobs that run in main """ cwd = os.getcwd() script = os.path.basename(__file__) js = JobScheduler(cwd + "/" + script, logDir="../../logs/", jobName="extractIntensity", partition='exfel',
def main(args): global device lowest_loss = 1e20 # create checkpoint dir and log dir if not isdir(args.checkpoint): print("Create new checkpoint folder " + args.checkpoint) mkdir_p(args.checkpoint) if not args.resume: if isdir(args.logdir): shutil.rmtree(args.logdir) mkdir_p(args.logdir) # create model model = models.__dict__["skinnet"](args.nearest_bone, args.Dg, args.Lf) model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # optionally resume from a checkpoint lr = args.lr if args.resume: if isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] lowest_loss = checkpoint['lowest_loss'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) lr = optimizer.param_groups[0]['lr'] print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) train_loader = DataLoader(SkinDataset(root=args.train_folder), batch_size=args.train_batch, shuffle=True) val_loader = DataLoader(SkinDataset(root=args.val_folder), batch_size=args.test_batch, shuffle=False) test_loader = DataLoader(SkinDataset(root=args.test_folder), batch_size=args.test_batch, shuffle=False) if args.evaluate: print('\nEvaluation only') test_loss = test(test_loader, model, args, save_result=True) print('test_loss {:6f}'.format(test_loss)) return scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, args.schedule, gamma=args.gamma) logger = SummaryWriter(log_dir=args.logdir) for epoch in range(args.start_epoch, args.epochs): lr = scheduler.get_last_lr() print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr[0])) train_loss = train(train_loader, model, optimizer, args) val_loss = test(val_loader, model, args) test_loss = test(test_loader, model, args) scheduler.step() print('Epoch{:d}. train_loss: {:.6f}.'.format(epoch + 1, train_loss)) print('Epoch{:d}. val_loss: {:.6f}.'.format(epoch + 1, val_loss)) print('Epoch{:d}. test_loss: {:.6f}.'.format(epoch + 1, test_loss)) # remember best acc and save checkpoint is_best = val_loss < lowest_loss lowest_loss = min(val_loss, lowest_loss) save_checkpoint({'epoch': epoch + 1, 'state_dict': model.state_dict(), 'lowest_loss': lowest_loss, 'optimizer': optimizer.state_dict()}, is_best, checkpoint=args.checkpoint) info = {'train_loss': train_loss, 'val_loss': val_loss, 'test_loss': test_loss} for tag, value in info.items(): logger.add_scalar(tag, value, epoch + 1)