def create_components(self, mode):
font = proj_dir("fonts", "Roboto-Regular.ttf")
self.model = self.create_model(None, mode)
self.logger = self.create_logger(proj_dir("logs"), font)
self.trainer = self.create_trainer(self.device, self.model, None, self.logger)
self.model_loader = self.create_model_loader(self.model_id)
def main():
device = "cuda:0"
parser = argparse.ArgumentParser()
parser.add_argument("--model", default=None)
parser.add_argument("--checkpoint", type=int, default=None)
parser.add_argument("--dataset", default=None)
parser.add_argument("--data_type", default="test")
args = parser.parse_args()
model_path = proj_dir("models", "ddff_mdff")
loader = ModelLoader(model_path, args.model)
model_params = loader.load("params")
model = project.create_component("net", **model_params["net"])
epoch = args.checkpoint
if args.checkpoint is not None:
print("Load from checkpoint")
model.load_state_dict(loader.load("checkpoint", epoch=epoch)["model"])
else:
print("Load model")
loader.load("model", model)
# data_class = datatype_from_str(args.dataset)
#data = data_class(root_dir=proj_dir("datasets"), data_type=args.data_type)
data = project.create_component("data",
name=args.dataset,
root_dir=proj_dir("datasets"),
data_type=args.data_type)
data_loader = torch.utils.data.DataLoader(data,
batch_size=1,
shuffle=False,
num_workers=1)
trainer = Trainer(model, device)
test_loss, eval_pred = trainer.test_model(data_loader)
loader.save("eval", {"pred": eval_pred},
epoch=epoch,
data_name=args.dataset,
data_type=args.data_type)
print("Done. Loss:", test_loss)
def train_model(params_file, model_name):
with open(params_file) as json_data:
params = json.load(json_data)
params = BatchTrain.set_param_paths(params,
dataset_path=proj_dir("datasets"),
model_path=proj_dir(
"models", "ddff_mdff"))
if model_name is not None:
params["save"]["name"] = model_name
batch_train = BatchTrain()
batch_train.train(params)
def simple_test_models(test_name,
model_setup,
model_id,
dataset_name,
ramps_per_clip,
ramp_length=4,
epoch=None,
colormap=None,
model_out_idx=-1):
model_desc = model_setup + "_testset"
setup = get_train_setup(model_setup, model_id, model_desc)
imgs = setup.simple_test(epoch,
model_id,
dataset_name,
ramp_length,
ramps_per_clip,
model_out_idx=model_out_idx)
path = proj_dir("logs", "simple_test",
test_name) # + "_" + time.strftime("%Y%m%d-%H%M%S"))
num_clips = imgs.shape[0] // ramps_per_clip
for i in range(num_clips):
save_image_batch(imgs[i * ramps_per_clip:(i + 1) * ramps_per_clip],
os.path.join(path, f"seq{i}"),
normalize="clip",
colormap=colormap)
def create_dataset(data_type):
data = VideoDepthFocusData(proj_dir("datasets"), data_type, "dining_room")
data.configure(sample_count=2,
sample_skip=1,
depth_output_indices=1,
use_allinfocus=False)
return data
def main():
ratio = [80, 20]
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", default="art_scene")
args = parser.parse_args()
use_slurm_system()
data_dir = proj_dir("datasets", args.dataset)
assert os.path.exists(proj_dir("datasets"))
data_paths = glob.glob(os.path.join(data_dir, "train", "*"))
data_paths = [os.path.basename(p) for p in data_paths]
split_dict = split_data(data_paths, ratio, ["train", "val"])
with open(os.path.join(data_dir, "split.json"), "w") as f:
json.dump(split_dict, f, indent=4)
def create_model(self, num_in_channels, mode):
assert num_in_channels == 3
return DenseNet3DAutoEncoder(
proj_dir("pretrained"),
use_concat=self.use_concat,
use_transp_conv=self.use_transp_conv,
decoder_conv_kernel_sizes=self.decoder_conv_kernel_sizes,
use_2d_dec=self.use_2d_dec,
depth_for_all=self.depth_for_all,
load_pretrained=self.load_pretrained)
def load_pretrained(model, model_type, train_setup=None, freeze=True):
if model_type == "flow_nvidia":
from net.flownet2 import load_flownet_nvidia
load_flownet_nvidia(model)
else:
name = get_saved_model_name(model, model_type, train_setup)
checkpoint = ModelLoader(proj_dir("pretrained"),
name,
require_exist=True).load("checkpoint")
model.load_state_dict(checkpoint["model"])
if freeze:
for param in model.parameters():
param.requires_grad = False
def download_data(link_file):
with open(link_file, "r") as f:
link_dict = json.load(f)
for k, v in link_dict.items():
print(k, len(v))
print("Total tex count", sum(len(links) for links in link_dict.values()))
base_url = "https://texture-ninja.nyc3.cdn.digitaloceanspaces.com/"
for category, links in link_dict.items():
for link in links:
filepath = proj_dir("crawl", "texninja", category, link)
download(filepath, base_url + link)
def pytorch_to_caffe(model_id, setup, epoch=None):
model = get_model(model_id, setup, epoch)
model.train(False)
x = torch.randn(1, 4, 3, 256, 256, requires_grad=True, device="cuda")
model_export_root = proj_dir("export", f"m{model_id}_{setup}")
os.makedirs(model_export_root)
model_export_path = os.path.join(model_export_root, "model.onnx")
# Export the model
torch_out = torch.onnx._export(model,
x,
model_export_path,
export_params=True)
prepared_backend = check_model(model_export_path, x, torch_out)
export_to_mobile(prepared_backend, model_export_root)
def main():
use_slurm_system()
url = "http://www.grsites.com"
url_tex = url + "/archive/textures"
home = get_tree(url_tex)
pages = get_pages(home)
global_idx = 0
for category, links in pages:
print(category)
for page_id, link in enumerate(links):
img_links = get_img_links(get_tree(url + link))
for i, img_link in enumerate(img_links):
filepath = proj_dir("crawl", "grsites", category, img_link.split("/")[-1])
download(filepath, img_link)
print(global_idx, category, page_id+1, "/", len(links), len(img_links), img_link, ">", filepath)
global_idx += 1
def get_dataset_path(self):
return proj_dir("datasets")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--update", action="store_true")
args = parser.parse_args()
VideoDepthFocusData.use_config = False
data = VideoDepthFocusData(proj_dir("datasets"), "train", "dining_room")
stats_per_frame = []
stats_per_clip = []
params = None
#args.update = True
if not args.update:
for clip in range(len(data)):
clip_stats = []
params = data._load_params(clip)
for idx in range(25):
depth = np.array(data._load_depth(clip, idx, params))
flow = np.array(data._load_flow(clip, idx, params, [512, 512]))
depth = depth[depth < 0xffff]
clip_stats.append({
"idx": idx,
"depth": {
"min": depth.min(),
"avg": depth.mean(),
"max": depth.max()
},
"flow": {
"min": flow.min(),
"avg": flow.mean(),
"max": flow.max(),
"nonzero": np.count_nonzero(flow)
}
})
clip_name = os.path.basename(data.clip_dirs[clip])
stats_per_frame.append({"clip": clip_name, "frames": clip_stats})
print("Clip", clip, "done")
with open(os.path.join(data.data_path, "stats_per_frame.json"),
"w") as f:
json.dump(stats_per_frame, f, default=default, indent=4)
else:
with open(os.path.join(data.data_path, "stats_per_frame.json"),
"r") as f:
stats_per_frame = json.load(f)
for clip in range(len(data)):
params = data._load_params(clip)
clip_name = os.path.basename(data.clip_dirs[clip])
stats_per_clip.append({
"clip:":
clip_name,
"frames":
dict_stat_compose(stats_per_frame[clip]["frames"])
})
stats_total = {
"focus_min": params["frames"][0]["focDist"],
"focus_max": params["frames"][-1]["focDist"],
**dict_stat_compose(stats_per_clip)
}
#print(stats_per_frame)
print(stats_per_clip)
print(stats_total)
with open(os.path.join(data.data_path, "stats_total.json"), "w") as f:
json.dump(stats_total, f, default=default, indent=4)
with open(os.path.join(data.data_path, "stats_per_clip.json"), "w") as f:
json.dump(stats_per_clip, f, default=default, indent=4)
from tools.project import proj_dir
# For parsing commandline arguments
parser = argparse.ArgumentParser()
parser.add_argument("--ffmpeg",
type=str,
default="ffmpeg",
help='path to ffmpeg')
parser.add_argument(
"--dataset",
type=str,
default="adobe240fps", # default="custom",
help='specify if using "adobe240fps" or custom video dataset')
parser.add_argument("--videos_folder",
type=str,
default=proj_dir("datasets", "adobe240fps", "raw"),
help='path to the folder containing videos')
parser.add_argument("--dataset_folder",
type=str,
default=proj_dir("datasets", "adobe240fps"),
help='path to the output dataset folder')
parser.add_argument("--img_width",
type=int,
default=640,
help="output image width")
parser.add_argument("--img_height",
type=int,
default=360,
help="output image height")
parser.add_argument("--train_test_split",
type=tuple,
def get_model_path():
return proj_dir("models")
focalstack_paths = \
[os.path.join(data_dir, focalstack_folder, image_id, "{}.jpg".format(i))
for i in range(focalstack_size)]
depth_path = \
os.path.join(data_dir, depth_folder, image_id, depthmap_name)
depth_img = Image.open(
os.path.join(root_dir, data_folder_name, depth_path))
to_float_tensor = data_transforms.ToFloatTensor()
val_crops = data_transforms.RandomCrop.get_all_valid_crop_pos(
to_float_tensor(depth_img), VideoDepthFocusData.crop_size)
print("Found {} crops".format(len(val_crops)))
data.append(focalstack_paths + [depth_path])
val_crops_list.append(val_crops)
tools.save_file(os.path.join(out_dir, "{}_data.csv".format(data_type)),
data)
crops_file = os.path.join(out_dir,
"{}_data_valid_crops.pkl".format(data_type))
with open(crops_file, "wb") as f:
pickle.dump(val_crops_list, f)
if __name__ == "__main__":
setup_data(proj_dir("datasets"), "mdff_setup")
def load_flownet_nvidia(model):
from tools.project import proj_dir
path = proj_dir("pretrained", "FlowNet2_checkpoint.pth.tar")
ckeckpoint = torch.load(path)
model.load_state_dict(ckeckpoint["state_dict"])
def test_model_set(test_name,
model_setup,
dataset_name,
loss="mse",
mask_img=False,
colormap=None,
model_labels=None,
append_cmap=False,
whole_seq_out=None,
img_format="jpg",
img_grid_nrow=8,
fixed_frame_indices=None,
select_focus_dists=None,
select_rel_indices=None,
normalize_tensors=False,
single_tensor_out=False,
**kwargs):
loggers = []
epochs = []
for i, setup_cfg in enumerate(model_setup):
(mid, setup_name) = parse_setup(setup_cfg)
if fixed_frame_indices is not None and isinstance(
fixed_frame_indices[0], list):
ffi = fixed_frame_indices[i]
else:
ffi = fixed_frame_indices
if select_focus_dists is not None and isinstance(
select_focus_dists[0], list):
sfd = select_focus_dists[i]
else:
sfd = select_focus_dists
setup = test_model(model_id=mid,
setup_name=setup_name,
dataset_name=dataset_name,
whole_seq_out=whole_seq_out,
fixed_frame_indices=ffi,
select_rel_indices=select_rel_indices,
select_focus_dists=sfd,
**kwargs)
loggers.append(setup.logger)
epochs.append(setup.trainer.epoch - 1)
mask_img_func = (lambda x: x <= 1) if mask_img else None
# path = proj_dir("logs", "test", time.strftime("%Y%m%d-%H%M%S") + "_" + dataset_name)
path = proj_dir("logs", "test",
test_name + "_" + time.strftime("%Y%m%d-%H%M%S"))
TrainLogger.save_all_test_images(
loggers,
mask_img_func=mask_img_func,
loss_label=loss,
model_labels=model_labels,
path=os.path.join(path, "img"),
colormap=colormap,
append_cmap=append_cmap,
whole_seq_out=whole_seq_out,
img_format=img_format,
epochs=epochs,
nrow=img_grid_nrow,
normalize_tensors=normalize_tensors,
single_tensor_out=single_tensor_out,
rotate_suwa=dataset_name == "suwajanakorn")
TrainLogger.save_all_test_stats(loggers, path=path)
return path
