def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED model = models.get_face_alignment_net(config) # copy model files writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } gpus = list(config.GPUS) model = nn.DataParallel(model, device_ids=gpus).cuda() # loss criterion = torch.nn.MSELoss(size_average=True).cuda() optimizer = utils.get_optimizer(config, model) best_nme = 100 last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'latest.pth') if os.path.islink(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_nme = checkpoint['best_nme'] model.module.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) else: print("=> no checkpoint found") if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) dataset_type = get_dataset(config) train_loader = DataLoader(dataset=dataset_type(config, is_train=True), batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) val_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() function.train(config, train_loader, model, criterion, optimizer, epoch, writer_dict) # evaluate nme, predictions = function.validate(config, val_loader, model, criterion, epoch, writer_dict) is_best = nme < best_nme best_nme = min(nme, best_nme) logger.info('=> saving checkpoint to {}'.format(final_output_dir)) print("best:", is_best) utils.save_checkpoint( { "state_dict": model.module.state_dict(), "epoch": epoch + 1, "best_nme": best_nme, "optimizer": optimizer.state_dict(), }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch)) final_model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn.benchmark = config.CUDNN.BENCHMARK # cudnn.determinstic = config.CUDNN.DETERMINISTIC # cudnn.enabled = config.CUDNN.ENABLED # if isinstance(config.TRAIN.LR_STEP, list): # lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( # optimizer, config.TRAIN.LR_STEP, # # config.TRAIN.LR_FACTOR, last_epoch-1 # config.TRAIN.LR_FACTOR, 0 # ) # else: # lr_scheduler = torch.optim.lr_scheduler.StepLR( # optimizer, config.TRAIN.LR_STEP, # # config.TRAIN.LR_FACTOR, last_epoch-1 # config.TRAIN.LR_FACTOR, 0 # ) dataset_type = get_dataset(config) train_dataset = dataset_type(config, is_train=True) # train_dataset[0] # return 0 train_loader = DataLoader( dataset=dataset_type(config, is_train=True), # batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus), batch_size=config.TRAIN.BATCH_SIZE_PER_GPU, shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS) # val_loader = DataLoader( # dataset=dataset_type(config, # is_train=True), # # batch_size=config.TEST.BATCH_SIZE_PER_GPU*len(gpus), # batch_size=config.TEST.BATCH_SIZE_PER_GPU, # shuffle=False, # num_workers=config.WORKERS, # # pin_memory=config.PIN_MEMORY # ) model = models.get_face_alignment_net(config) # copy model files writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } gpus = list(config.GPUS) # model = nn.DataParallel(model, device_ids=gpus).cuda() model.to("cuda") # loss criterion = torch.nn.MSELoss(size_average=True).cuda() # criterion = fnn.mse_loss # criterion = WingLoss() # criterion = Loss_weighted() optimizer = utils.get_optimizer(config, model) best_nme = 100 last_epoch = config.TRAIN.BEGIN_EPOCH if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'latest.pth') if os.path.isfile(model_state_file): with open(model_state_file, "rb") as fp: state_dict = torch.load(fp) model.load_state_dict(state_dict) last_epoch = 1 # checkpoint = torch.load(model_state_file) # last_epoch = checkpoint['epoch'] # best_nme = checkpoint['best_nme'] # model.load_state_dict(checkpoint['state_dict']) # optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint (epoch {})" .format(last_epoch)) else: print("=> no checkpoint found") lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() function.train(config, train_loader, model, criterion, optimizer, epoch, writer_dict) # evaluate nme = 0 # nme, predictions = function.validate(config, val_loader, model, # criterion, epoch, writer_dict) is_best = True # is_best = nme < best_nme best_nme = min(nme, best_nme) logger.info('=> saving checkpoint to {}'.format(final_output_dir)) print("best:", is_best) torch.save(model.state_dict(), os.path.join(final_output_dir, 'mse_relu_checkpoint_{}.pth'.format(epoch))) # utils.save_checkpoint( # {"state_dict": model, # "epoch": epoch + 1, # "best_nme": best_nme, # "optimizer": optimizer.state_dict(), # }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch)) final_model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info('saving final model state to {}'.format( final_model_state_file)) torch.save(model.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): # args = parse_args() # logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') # logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) # gpus = list(config.GPUS) # # # load model state_dict = torch.load(args.model_file) if 'state_dict' in state_dict.keys(): state_dict = state_dict['state_dict'] model.load_state_dict(state_dict) else: model.module.load_state_dict(state_dict) model = nn.DataParallel(model, device_ids=gpus).cuda() # dataset_type = get_dataset(config) test_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) nme, predictions = function.inference(config, test_loader, model) torch.save(predictions, os.path.join(final_output_dir, 'predictions.pth')) target = test_loader.dataset.load_all_pts() pred = 16 * predictions l = len(pred) res = 0.0 res_tmp = [0.0 for i in range(config.MODEL.NUM_JOINTS)] res_tmp = np.array(res_tmp) res_temp_x = target - pred res_temp_x = res_temp_x[:, :, 0] res_temp_y = target - pred res_temp_y = res_temp_y[:, :, 1] # csv_file_test_x = pd.DataFrame(np.transpose(np.array(pred[:, :, 0])), columns=test_loader.dataset.annotation_files) # csv_file_test_y = pd.DataFrame(np.transpose(np.array(pred[:, :, 1])), columns=test_loader.dataset.annotation_files) # csv_file_target_x = pd.DataFrame(np.transpose(np.array(target[:, :, 0])), columns=test_loader.dataset.annotation_files) # csv_file_target_y = pd.DataFrame(np.transpose(np.array(target[:, :, 1])), columns=test_loader.dataset.annotation_files) for i in range(l): trans = np.sqrt( pow(target[i][0][0] - target[i][1][0], 2) + pow(target[i][0][1] - target[i][1][1], 2)) / 30.0 res_temp_x[i] = res_temp_x[i] / trans res_temp_y[i] = res_temp_y[i] / trans for j in range(len(target[i])): dist = np.sqrt( np.power((target[i][j][0] - pred[i][j][0]), 2) + np.power((target[i][j][1] - pred[i][j][1]), 2)) / trans res += dist res_tmp[j] += dist res_t = np.sqrt(res_temp_x * res_temp_x + res_temp_y * res_temp_y) # pd.DataFrame(data=res_temp_x.data.value).to_csv('res_x') # pd.DataFrame(data=res_temp_y.data.value).to_csv('res_y') # pd.DataFrame(data=res_t.data.value).to_csv('res_t') res_tmp /= np.float(len(pred)) print(res_tmp) print(np.mean(res_tmp)) res /= (len(pred) * len(pred[0])) print(res)
def main(): # parsing script arguments args = parse_args() device = torch.device(args.device) # initialize logger logger, final_output_dir, tb_log_dir = utils.create_logger(config, args.cfg, "demo") # log arguments and config values logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # init landmark model model = models.get_face_alignment_net(config) # get input size from the config input_size = config.MODEL.IMAGE_SIZE # load model state_dict = torch.load(args.landmark_model, map_location=device) # remove `module.` prefix from the pre-trained weights new_state_dict = OrderedDict() for key, value in state_dict.items(): name = key[7:] new_state_dict[name] = value # load weights without the prefix model.load_state_dict(new_state_dict) # run model on device model = model.to(device) # init mean and std values for the landmark model's input mean = config.MODEL.MEAN mean = np.array(mean, dtype=np.float32) std = config.MODEL.STD std = np.array(std, dtype=np.float32) # defining prototxt and caffemodel paths detector_model = args.detector_model detector_weights = args.detector_weights # load model detector = cv2.dnn.readNetFromCaffe(detector_model, detector_weights) capture = cv2.VideoCapture(0) frame_num = 0 while True: # capture frame-by-frame success, frame = capture.read() # break if no frame if not success: break frame_num += 1 print("frame_num: ", frame_num) landmarks_img = frame.copy() result = frame.copy() result = result.astype(np.float32) / 255.0 # get frame's height and width height, width = frame.shape[:2] # 640x480 # resize and subtract BGR mean values, since Caffe uses BGR images for input blob = cv2.dnn.blobFromImage( frame, scalefactor=1.0, size=(300, 300), mean=(104.0, 177.0, 123.0), ) # passing blob through the network to detect faces detector.setInput(blob) # detector output format: # [image_id, class, confidence, left, bottom, right, top] face_detections = detector.forward() # loop over the detections for i in range(0, face_detections.shape[2]): # extract confidence confidence = face_detections[0, 0, i, 2] # filter detections by confidence greater than the minimum threshold if confidence > 0.5: # get coordinates of the bounding box box = face_detections[0, 0, i, 3:7] * np.array( [width, height, width, height], ) (x1, y1, x2, y2) = box.astype("int") # show original image cv2.imshow("original image", frame) # crop to detection and resize resized = crop( frame, torch.Tensor([x1 + (x2 - x1) / 2, y1 + (y2 - y1) / 2]), 1.5, tuple(input_size), ) # convert from BGR to RGB since HRNet expects RGB format resized = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) img = resized.astype(np.float32) / 255.0 # normalize landmark net input normalized_img = (img - mean) / std # predict face landmarks model = model.eval() with torch.no_grad(): input = torch.Tensor(normalized_img.transpose([2, 0, 1])) input = input.to(device) output = model(input.unsqueeze(0)) score_map = output.data.cpu() preds = decode_preds( score_map, [torch.Tensor([x1 + (x2 - x1) / 2, y1 + (y2 - y1) / 2])], [1.5], score_map.shape[2:4], ) preds = preds.squeeze(0) landmarks = preds.data.cpu().detach().numpy() # draw landmarks for k, landmark in enumerate(landmarks, 1): landmarks_img = cv2.circle( landmarks_img, center=(landmark[0], landmark[1]), radius=3, color=(0, 0, 255), thickness=-1, ) # draw landmarks' labels landmarks_img = cv2.putText( img=landmarks_img, text=str(k), org=(int(landmark[0]) + 5, int(landmark[1]) + 5), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0, 0, 255), ) # show results by drawing predicted landmarks and their labels cv2.imshow("image with landmarks", landmarks_img) # get chosen landmarks 2-16, 30 as destination points # note that landmarks numbering starts from 0 dst_pts = np.array( [ landmarks[1], landmarks[2], landmarks[3], landmarks[4], landmarks[5], landmarks[6], landmarks[7], landmarks[8], landmarks[9], landmarks[10], landmarks[11], landmarks[12], landmarks[13], landmarks[14], landmarks[15], landmarks[29], ], dtype="float32", ) # load mask annotations from csv file to source points mask_annotation = os.path.splitext(os.path.basename(args.mask_image))[0] mask_annotation = os.path.join( os.path.dirname(args.mask_image), mask_annotation + ".csv", ) with open(mask_annotation) as csv_file: csv_reader = csv.reader(csv_file, delimiter=",") src_pts = [] for i, row in enumerate(csv_reader): # skip head or empty line if it's there try: src_pts.append(np.array([float(row[1]), float(row[2])])) except ValueError: continue src_pts = np.array(src_pts, dtype="float32") # overlay with a mask only if all landmarks have positive coordinates: if (landmarks > 0).all(): # load mask image mask_img = cv2.imread(args.mask_image, cv2.IMREAD_UNCHANGED) mask_img = mask_img.astype(np.float32) mask_img = mask_img / 255.0 # get the perspective transformation matrix M, _ = cv2.findHomography(src_pts, dst_pts) # transformed masked image transformed_mask = cv2.warpPerspective( mask_img, M, (result.shape[1], result.shape[0]), None, cv2.INTER_LINEAR, cv2.BORDER_CONSTANT, ) # mask overlay alpha_mask = transformed_mask[:, :, 3] alpha_image = 1.0 - alpha_mask for c in range(0, 3): result[:, :, c] = ( alpha_mask * transformed_mask[:, :, c] + alpha_image * result[:, :, c] ) # display the resulting frame cv2.imshow("image with mask overlay", result) # waiting for the escape button to exit k = cv2.waitKey(1) if k == 27: break # when everything done, release the capture capture.release() cv2.destroyAllWindows()
def main(): args = parse_args() # set logger and dir logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.experiment_name, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # set cudnn cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED # 目前仅支持单gpu,todo:增加多gpu支持 # set model and loss and criterion model = models.get_face_alignment_net(config) model = model.cuda(config.GPUS[0]) criterion = torch.nn.MSELoss(size_average=True).cuda(config.GPUS[0]) # criterion = AdaptiveWingLoss() optimizer = utils.get_optimizer(config, model) # get dataset dataset_type = get_dataset(config) # get dataloader train_loader = DataLoader(dataset=dataset_type(config, is_train=True), batch_size=config.TRAIN.BATCH_SIZE_PER_GPU, shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) val_loader = DataLoader(dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU, shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) # set lr_scheduler last_epoch = config.TRAIN.BEGIN_EPOCH if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) # set training writer writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } # set training resume function if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'latest.pth') if os.path.islink(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_nme = checkpoint['best_nme'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) else: print("=> no checkpoint found") # starting training best_nme = 10000 for epoch in range(last_epoch, config.TRAIN.END_EPOCH): lr_scheduler.step() # traing function.train(config, train_loader, model, criterion, optimizer, epoch, writer_dict) # evaluating nme, predictions = function.validate(config, val_loader, model, criterion, epoch, writer_dict) # saving is_best = nme < best_nme best_nme = min(nme, best_nme) logger.info('=> saving checkpoint to {}'.format(final_output_dir)) print("best:", is_best) utils.save_checkpoint( { "state_dict": model, "epoch": epoch + 1, "best_nme": best_nme, "optimizer": optimizer.state_dict(), }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch)) final_model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) # cudnn.benchmark = config.CUDNN.BENCHMARK # cudnn.determinstic = config.CUDNN.DETERMINISTIC # cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) gpus = list(config.GPUS) # model = nn.DataParallel(model, device_ids=gpus).cuda() model.to("cuda") # print(model) # load model # state_dict = torch.load(args.model_file) # print(state_dict) # model = torch.load(args.model_file) with open(args.model_file, "rb") as fp: state_dict = torch.load(fp) model.load_state_dict(state_dict) # model.load_state_dict(state_dict['state_dict']) # if 'state_dict' in state_dict.keys(): # state_dict = state_dict['state_dict'] # # print(state_dict) # model.load_state_dict(state_dict) # else: # model.module.load_state_dict(state_dict) dataset_type = get_dataset(config) test_loader = DataLoader( dataset=dataset_type(config, is_train=False), batch_size=config.TEST.BATCH_SIZE_PER_GPU*len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY ) predictions = function.inference(config, test_loader, model) # print("len(predictions)", len(predictions)) # print(predictions[0]) df_predictions = [] for pred in predictions: row = dict() row['file_name'] = pred[0] for id_point in range(194): row[f'Point_M{id_point}_X'] = int(pred[1][id_point]) row[f'Point_M{id_point}_Y'] = int(pred[2][id_point]) df_predictions.append(row) df_predictions = pd.DataFrame(df_predictions) # print(predictions_meta[0]) df_predictions.to_csv('pred_test.csv', index=False)
def main(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'train') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED gpus = list(config.GPUS) dataset_type = get_dataset(config) train_data = dataset_type(config, split="train") train_loader = DataLoader(dataset=train_data, batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus), shuffle=config.TRAIN.SHUFFLE, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) val_data = dataset_type(config, split="valid") val_loader = DataLoader(dataset=val_data, batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus), shuffle=False, num_workers=config.WORKERS, pin_memory=config.PIN_MEMORY) # config.MODEL.NUM_JOINTS = train_data.get_num_points() model = models.get_face_alignment_net(config) # copy model files writer_dict = { 'writer': SummaryWriter(log_dir=tb_log_dir), 'train_global_steps': 0, 'valid_global_steps': 0, } model = nn.DataParallel(model, device_ids=gpus).cuda() # loss criterion = torch.nn.MSELoss(size_average=True).cuda() optimizer = utils.get_optimizer(config, model) best_nme = 100 last_epoch = config.TRAIN.BEGIN_EPOCH if isinstance(config.TRAIN.LR_STEP, list): lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) else: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR, last_epoch - 1) if config.TRAIN.RESUME: model_state_file = os.path.join(final_output_dir, 'final.pth') if os.path.islink(model_state_file): checkpoint = torch.load(model_state_file) last_epoch = checkpoint['epoch'] best_nme = checkpoint['best_nme'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) else: print("=> no checkpoint found") loss = [] for epoch in range(last_epoch, config.TRAIN.END_EPOCH): losses, diff = function.train(config, train_loader, model, criterion, optimizer, epoch, writer_dict) loss.append(losses) lr_scheduler.step() np.save( os.path.join(final_output_dir, "train_diff@epoch{}".format(epoch)), diff) # evaluate nme, predictions, diff = function.validate(config, val_loader, model, criterion, epoch, writer_dict) np.save( os.path.join(final_output_dir, "valid_diff@epoch{}".format(epoch)), diff) is_best = nme < best_nme best_nme = min(nme, best_nme) logger.info('=> saving checkpoint to {}'.format(final_output_dir)) print("best:", is_best) utils.save_checkpoint( { "state_dict": model, "epoch": epoch + 1, "best_nme": best_nme, "optimizer": optimizer.state_dict(), }, predictions, is_best, final_output_dir, 'checkpoint_{}.pth'.format(epoch)) if is_best: for i in range(len(predictions)): afile = val_data.annotation_files[i] new_afile = '{}.{}.txt'.format( afile, os.path.basename(args.cfg).split('.')[0]) with open(new_afile, 'wt') as f: pts = predictions[i].cpu().numpy() for j in range(len(pts)): f.write("{},{}\n".format( pts[j][1] / val_data.factor[1], pts[j][0] / val_data.factor[0])) pd.DataFrame(data=loss).to_csv('loss2.csv') final_model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info( 'saving final model state to {}'.format(final_model_state_file)) torch.save(model.module.state_dict(), final_model_state_file) writer_dict['writer'].close()
def main_wflw(): args = parse_args() logger, final_output_dir, tb_log_dir = \ utils.create_logger(config, args.cfg, 'test') logger.info(pprint.pformat(args)) logger.info(pprint.pformat(config)) cudnn.benchmark = config.CUDNN.BENCHMARK cudnn.determinstic = config.CUDNN.DETERMINISTIC cudnn.enabled = config.CUDNN.ENABLED config.defrost() config.MODEL.INIT_WEIGHTS = False config.freeze() model = models.get_face_alignment_net(config) dataset_type = get_dataset(config) dataset = dataset_type(config, is_train=True) fp = open('data/wflw/face_landmarks70_wflw_train.csv', 'w') for i in range(len(dataset)): img, image_path, meta = dataset[i] fold, name = image_path.split('/')[-2], image_path.split('/')[-1] folder = osp.join('data/wflw/xximages', fold) if not osp.exists(folder): os.makedirs(folder) fname = osp.join(folder, name) scale = meta['scale'] center = meta['center'] tpts = meta['tpts'] selpts = [] for j in range(0, 33, 2): selpts.append(tpts[j]) # eyebow selpts.append(tpts[33]) selpts.append((tpts[34] + tpts[41]) / 2) selpts.append((tpts[35] + tpts[40]) / 2) selpts.append((tpts[36] + tpts[39]) / 2) selpts.append((tpts[37] + tpts[38]) / 2) selpts.append((tpts[42] + tpts[50]) / 2) selpts.append((tpts[43] + tpts[49]) / 2) selpts.append((tpts[44] + tpts[48]) / 2) selpts.append((tpts[45] + tpts[47]) / 2) selpts.append(tpts[46]) # nose for j in range(51, 60): selpts.append(tpts[j]) # eye selpts.append(tpts[60]) selpts.append((tpts[61] + tpts[62]) / 2) selpts.append(tpts[63]) selpts.append(tpts[64]) selpts.append(tpts[65]) selpts.append((tpts[66] + tpts[67]) / 2) selpts.append(tpts[68]) selpts.append(tpts[69]) selpts.append((tpts[70] + tpts[71]) / 2) selpts.append(tpts[72]) selpts.append((tpts[73] + tpts[74]) / 2) selpts.append(tpts[75]) for j in range(76, 98): selpts.append(tpts[j]) fp.write('%s,%.2f,%.1f,%.1f' % (osp.join(fold, name), scale, center[0], center[1])) for spt in selpts: cv2.circle(img, (spt[0], spt[1]), 1, (0, 0, 255)) fp.write(',%f,%f' % (spt[0], spt[1])) fp.write('\n') img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) cv2.imwrite(fname, img) fp.close()