if args.dataset_type == "Avenue_rain":
for single_dir in train_folder:
if not os.path.exists(single_dir):
rain_type = single_dir.strip().split('frames/')[1].strip().split(
"_")[0]
print(
"creating training images with augmentation %s and rain level 0.70"
% rain_type)
aug_data.save_avenue_rain_or_bright(args.dataset_path,
rain_type,
True,
"training",
bright_space=0.7)
frame_trans = data_utils.give_frame_trans(args.dataset_type, [args.h, args.w])
train_dataset = data_utils.DataLoader(train_folder,
frame_trans,
time_step=args.t_length - 1,
num_pred=1)
test_dataset = data_utils.DataLoader(test_folder,
frame_trans,
time_step=args.t_length - 1,
num_pred=1)
train_batch = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
drop_last=True)
args = parser.parse_args()
torch.manual_seed(2020)
torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# the path of each dataset
s_train, s_test = data_utils.give_data_folder(args.source_dataset,
args.dataset_path)
t_train, t_test = data_utils.give_data_folder(args.target_dataset,
args.dataset_path)
# prepare image transform
s_frame_trans = data_utils.give_frame_trans(args.source_dataset,
[args.h, args.w])
t_frame_trans = data_utils.give_frame_trans(args.target_dataset,
[args.h, args.w])
# prepare dataset
# s_test_label = np.load(args.source_test_label_path, allow_pickle=True)
t_test_label = np.load(args.target_test_label_path, allow_pickle=True)
s_train_dataset = data_utils.DataLoader(s_train,
s_frame_trans,
None,
True,
time_step=args.t_length - 1,
num_pred=1,
video_start=1,
video_end=5)