GIMP-ML/gimp-plugins/MiDaS/run.py
2020-07-17 09:30:16 +05:30

79 lines
2.1 KiB
Python

"""Compute depth maps for images in the input folder.
"""
# import os
# import glob
import torch
# from monodepth_net import MonoDepthNet
# import utils
# import matplotlib.pyplot as plt
import numpy as np
import cv2
# import imageio
def run_depth(img, model_path, Net, utils, target_w=None):
"""Run MonoDepthNN to compute depth maps.
Args:
input_path (str): path to input folder
output_path (str): path to output folder
model_path (str): path to saved model
"""
# print("initialize")
# select device
device = torch.device("cpu")
# print("device: %s" % device)
# load network
model = Net(model_path)
model.to(device)
model.eval()
# get input
# img_names = glob.glob(os.path.join(input_path, "*"))
# num_images = len(img_names)
# create output folder
# os.makedirs(output_path, exist_ok=True)
# print("start processing")
# for ind, img_name in enumerate(img_names):
# print(" processing {} ({}/{})".format(img_name, ind + 1, num_images))
# input
# img = utils.read_image(img_name)
w = img.shape[1]
scale = 640. / max(img.shape[0], img.shape[1])
target_height, target_width = int(round(img.shape[0] * scale)), int(round(img.shape[1] * scale))
img_input = utils.resize_image(img)
# print(img_input.shape)
img_input = img_input.to(device)
# compute
with torch.no_grad():
out = model.forward(img_input)
depth = utils.resize_depth(out, target_width, target_height)
img = cv2.resize((img * 255).astype(np.uint8), (target_width, target_height), interpolation=cv2.INTER_AREA)
# np.save(filename + '.npy', depth)
# utils.write_depth(filename, depth, bits=2)
depth_min = depth.min()
depth_max = depth.max()
bits = 1
max_val = (2 ** (8 * bits)) - 1
if depth_max - depth_min > np.finfo("float").eps:
out = max_val * (depth - depth_min) / (depth_max - depth_min)
else:
out = 0
out = out.astype("uint8")
# cv2.imwrite("out.png", out)
return out
# print("finished")