def predict(config_file): """ Main function that runs predictions Args: config_file [str]: path to config file Returns: None """ ################## # configure logger ################## logger = set_logger("../log/predict.log") ################## # Load config from config file ################## logger.info(f"Load config from {config_file}") config = parse_config(config_file) image_width = config['common']['in_image_width'] image_height = config['common']['in_image_height'] predict_img = config['predict']['folder_path'] weights_path = config['common']['weights_path'] X,img_names = preprocess(predict_img, image_width, image_height) model = KeyPointModel().getModel() logger.info(f"Loading weights from {weights_path}") model.load_weights(weights_path) # logger.info("-----------Model Summary------------") # logger.info(model.summary()) predicted_keypoints = model.predict(X) logger.info("Prediction Completed. Writing output to predicted.csv") write_output(predicted_keypoints, img_names)
transforms_test = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((360, 480)), transforms.ToTensor(), transforms.Normalize(mean=[0.4372, 0.4372, 0.4373], std=[0.2479, 0.2475, 0.2485]) ]) datasets_test = KeyPointDatasets(root_dir="./data", transforms=transforms_test) dataloader_test = DataLoader(datasets_test, batch_size=4, shuffle=True, collate_fn=datasets_test.collect_fn) model = KeyPointModel() model.load_state_dict(torch.load("weights/epoch_290_0.232.pt")) img_list = glob.glob(os.path.join("./data/images", "*.jpg")) save_path = "./output" img_tensor_list = [] img_name_list = [] for i in range(len(img_list)): img_path = img_list[i] img_name = os.path.basename(img_path) img_name_list.append(img_name)
def train(config_file): """ Main function that train and persists model based on training set/ Args: config_file [str]: path to config file Returns: None """ ################ # config logger ################ logger = set_logger("../log/train.log") ############################### # Load config from config file ############################### logger.info(f"Load config from {config_file}") config = parse_config(config_file) keypoints_csv = Path(config['common']['labels_csv_path']) val_split = config['common']['val_split'] train_img_scr_path = config['common']['img_source_path'] test_img_scr_path = config['common']['img_source_path'] image_width = config['common']['in_image_width'] image_height = config['common']['in_image_height'] epochs = config['train']['epochs'] train_batch_size = config['train']['batch_size'] weight_path = config['common']['weight_path'] no_of_aug = config['train']['no_of_aug'] test_batch_size = config['test']['batch_size'] ############ # Load Data ############ logger.info(f"----------------Load the data----------------") selected_img, keypoint_df = load_data(keypoints_csv) logger.info(f"Number of selected images are {selected_img.shape}") logger.info(f"Few of the selected images are {selected_img[0:5]}") #################################### # Get train and test data generators #################################### X_train, y_train, X_test, y_test = train_test_split( selected_img, keypoint_df, val_split) train_gen = Car(x_set=X_train, y_set=y_train, mode='Train', data_path=train_img_scr_path, image_width=image_width, image_height=image_height, batch_size=train_batch_size, augmentations='Self', no_of_aug=no_of_aug) test_gen = Car( x_set=X_test, y_set=y_test, mode='Test', data_path=test_img_scr_path, image_width=image_width, image_height=image_height, batch_size=test_batch_size, ) ##################### # Set and train model ##################### logger.info( f"-------------------------Initiate Model---------------------") model = KeyPointModel().getModel() logger.info( f"--------------------Model Summary---------------------------") logger.info(f"{model.summary}") # compile the model model.compile(loss='mean_squared_error', optimizer='adam', metrics=['mean_absolute_error']) # modelCheckPoint = ModelCheckpoint('car-{val_loss:.2f}.h5', monitor='val_loss', verbose=1, save_best_only=True, save_weights_only=True) earlyS = EarlyStopping(monitor='val_loss', min_delta=1, patience=3, restore_best_weights=True) reducelr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=2, min_lr=1e-7) history = model.fit(x=train_gen, validation_data=test_gen, callbacks=[earlyS, reducelr], epochs=epochs) logger.info(history) logger.info("------------Saving Weights--------------") model.save_weights(weight_path)
transforms_all = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((480, 360)), transforms.ToTensor(), transforms.Normalize(mean=[0.4372, 0.4372, 0.4373], std=[0.2479, 0.2475, 0.2485]) ]) datasets = KeyPointDatasets(root_dir="./data", transforms=transforms_all) data_loader = DataLoader(datasets, shuffle=True, batch_size=bs, collate_fn=datasets.collect_fn) model = KeyPointModel() optimizer = torch.optim.Adam(model.parameters(), lr=3e-4) # criterion = torch.nn.SmoothL1Loss() criterion = torch.nn.MSELoss() scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1) for epoch in range(total_epoch): train(model, epoch, data_loader, optimizer, criterion) loss = test(model, epoch, data_loader, criterion) if epoch % 10 == 0: torch.save(model.state_dict(), "weights/epoch_%d_%.3f.pt" % (epoch, loss * 1000))
transforms_all = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((360, 480)), transforms.ToTensor(), transforms.Normalize(mean=[0.4372, 0.4372, 0.4373], std=[0.2479, 0.2475, 0.2485]) ]) dataset = KeyPointDatasets(root_dir="./data", transforms=transforms_all) dataloader = DataLoader(dataset, shuffle=True, batch_size=1, collate_fn=dataset.collect_fn) model = KeyPointModel() model.load_state_dict(torch.load(args.model)) for iter, (image, label) in enumerate(dataloader): # print(image.shape) bs = image.shape[0] hm = model(image) hm = _nms(hm) scores, inds, clses, ys, xs = _topk(hm, K=1) print(scores, '\n', inds, '\n', clses, '\n', ys, '\n', xs) hm = hm.detach().numpy()