Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| import numpy as np | |
| from tqdm import trange | |
| from PIL import Image | |
| def get_state(gpu): | |
| import torch | |
| midas = torch.hub.load("intel-isl/MiDaS", "MiDaS") | |
| if gpu: | |
| midas.cuda() | |
| midas.eval() | |
| midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms") | |
| transform = midas_transforms.default_transform | |
| state = {"model": midas, | |
| "transform": transform} | |
| return state | |
| def depth_to_rgba(x): | |
| assert x.dtype == np.float32 | |
| assert len(x.shape) == 2 | |
| y = x.copy() | |
| y.dtype = np.uint8 | |
| y = y.reshape(x.shape+(4,)) | |
| return np.ascontiguousarray(y) | |
| def rgba_to_depth(x): | |
| assert x.dtype == np.uint8 | |
| assert len(x.shape) == 3 and x.shape[2] == 4 | |
| y = x.copy() | |
| y.dtype = np.float32 | |
| y = y.reshape(x.shape[:2]) | |
| return np.ascontiguousarray(y) | |
| def run(x, state): | |
| model = state["model"] | |
| transform = state["transform"] | |
| hw = x.shape[:2] | |
| with torch.no_grad(): | |
| prediction = model(transform((x + 1.0) * 127.5).cuda()) | |
| prediction = torch.nn.functional.interpolate( | |
| prediction.unsqueeze(1), | |
| size=hw, | |
| mode="bicubic", | |
| align_corners=False, | |
| ).squeeze() | |
| output = prediction.cpu().numpy() | |
| return output | |
| def get_filename(relpath, level=-2): | |
| # save class folder structure and filename: | |
| fn = relpath.split(os.sep)[level:] | |
| folder = fn[-2] | |
| file = fn[-1].split('.')[0] | |
| return folder, file | |
| def save_depth(dataset, path, debug=False): | |
| os.makedirs(path) | |
| N = len(dset) | |
| if debug: | |
| N = 10 | |
| state = get_state(gpu=True) | |
| for idx in trange(N, desc="Data"): | |
| ex = dataset[idx] | |
| image, relpath = ex["image"], ex["relpath"] | |
| folder, filename = get_filename(relpath) | |
| # prepare | |
| folderabspath = os.path.join(path, folder) | |
| os.makedirs(folderabspath, exist_ok=True) | |
| savepath = os.path.join(folderabspath, filename) | |
| # run model | |
| xout = run(image, state) | |
| I = depth_to_rgba(xout) | |
| Image.fromarray(I).save("{}.png".format(savepath)) | |
| if __name__ == "__main__": | |
| from taming.data.imagenet import ImageNetTrain, ImageNetValidation | |
| out = "data/imagenet_depth" | |
| if not os.path.exists(out): | |
| print("Please create a folder or symlink '{}' to extract depth data ".format(out) + | |
| "(be prepared that the output size will be larger than ImageNet itself).") | |
| exit(1) | |
| # go | |
| dset = ImageNetValidation() | |
| abspath = os.path.join(out, "val") | |
| if os.path.exists(abspath): | |
| print("{} exists - not doing anything.".format(abspath)) | |
| else: | |
| print("preparing {}".format(abspath)) | |
| save_depth(dset, abspath) | |
| print("done with validation split") | |
| dset = ImageNetTrain() | |
| abspath = os.path.join(out, "train") | |
| if os.path.exists(abspath): | |
| print("{} exists - not doing anything.".format(abspath)) | |
| else: | |
| print("preparing {}".format(abspath)) | |
| save_depth(dset, abspath) | |
| print("done with train split") | |
| print("done done.") | |