File size: 3,501 Bytes
d8eaf88 23e9852 61d09e6 10c55f7 61d09e6 fcda6f5 c372dd5 d8eaf88 faac476 e32d51c 604f38b d8eaf88 59fba04 61d09e6 c372dd5 61d09e6 2c0045e 59fba04 604f38b eac489d 360dad8 59fba04 9ff509b 491b981 8f32fb8 491b981 9ff509b 49af897 360dad8 604f38b 4edfa31 e32d51c 23e9852 cdcc96d cc6c676 98d2c0f 4b2bdc6 faac476 4b2bdc6 faac476 4b2bdc6 cdcc96d 0a1c1a2 4b2bdc6 d5307ab 604f38b d5307ab faac476 fb6c0ad cc6c676 cdcc96d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import os
import gradio as gr
import cv2
import torch
import urllib.request
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import subprocess
def calculate_depth(model_type, gan_type, dim, slider, img):
if not os.path.exists('temp'):
os.system('mkdir temp')
filename = "Images/Input-Test/1.png"
img.save(filename, "PNG")
midas = torch.hub.load("intel-isl/MiDaS", model_type)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
midas.to(device)
midas.eval()
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
transform = midas_transforms.dpt_transform
else:
transform = midas_transforms.small_transform
img = cv2.imread(filename)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
input_batch = transform(img).to(device)
with torch.no_grad():
prediction = midas(input_batch)
prediction = torch.nn.functional.interpolate(
prediction.unsqueeze(1),
size=img.shape[:2],
mode="bicubic",
align_corners=False,
).squeeze()
output = prediction.cpu().numpy()
formatted = (output * 255.0 / np.max(output)).astype('uint8')
out_im = Image.fromarray(formatted)
out_im.save("Images/Input-Test/1_d.png", "PNG")
c_images = '1'
name_output = 'out'
dict_saved_gans = {'Cycle': '74962_110', 'Cycle(half)': '66942','noCycle': '31219_110', 'noCycle-noCr': '92332_110', 'noCycle-noCr-noL1': '82122_110', 'OnlyGen': '70944_110' }
subprocess.run(["python", "main.py", "--gan_type", 'WiggleGAN', "--expandGen", "4", "--expandDis", "4", "--batch_size", c_images, "--cIm", c_images,
"--visdom", "false", "--wiggleDepth", str(slider), "--seedLoad", dict_saved_gans[gan_type], "--gpu_mode", "false", "--imageDim", dim, "--name_wiggle", name_output
])
subprocess.run(["python", "WiggleResults/split.py", "--dim", dim])
path_video = os.path.join(os.path.dirname(__file__), 'WiggleResults' , name_output + '_0.mp4')
print(path_video)
return [out_im,f'WiggleResults/' + name_output + '_0.gif', path_video, f'WiggleResults/'+ name_output + '.jpg']
with gr.Blocks() as demo:
gr.Markdown("Start typing below and then click **Run** to see the output.")
## Depth Estimation
midas_models = ["DPT_Large","DPT_Hybrid","MiDaS_small"]
gan_models = ["Cycle","Cycle(half)","noCycle","noCycle-noCr","noCycle-noCr-noL1","OnlyGen"]
dim = ['256','512','1024']
with gr.Row():
inp = [gr.inputs.Dropdown(midas_models, default="MiDaS_small", label="Depth estimation model type")]
inp.append(gr.inputs.Dropdown(gan_models, default="Cycle", label="Different GAN trainings"))
inp.append(gr.inputs.Dropdown(dim, default="256", label="Wiggle dimension result"))
inp.append(gr.Slider(1,15, default = 2, label='StepCycles',step= 1))
with gr.Row():
inp.append(gr.Image(type="pil", label="Input"))
out = [gr.Image(type="pil", label="depth_estimation")]
with gr.Row():
out.append(gr.Image(type="file", label="Output_wiggle_gif"))
out.append(gr.Video(label="Output_wiggle_video"))
out.append(gr.Image(type="file", label="Output_images"))
btn = gr.Button("Calculate depth + Wiggle")
btn.click(fn=calculate_depth, inputs=inp, outputs=out)
demo.launch() |