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