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, img): if not os.path.exists('temp'): os.system('mkdir temp') filename = "Images/Input-Test/0.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/0_d.png", "PNG") return out_im def wiggle_effect(slider): dim = '256' c_images = '1' 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", '31219_110', "--gpu_mode", "false", "--imageDim", dim, "--name_wiggle", "result" ]) subprocess.run(["python", "WiggleResults/split.py", "--dim", dim]) return [f'WiggleResults/result.png',f'WiggleResults/result_0.gif'] 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"] inp = [gr.inputs.Dropdown(midas_models, default="MiDaS_small", label="Depth estimation model type")] with gr.Row(): inp.append(gr.Image(type="pil", label="Input")) out = gr.Image(type="pil", label="depth_estimation") btn = gr.Button("Calculate depth") btn.click(fn=calculate_depth, inputs=inp, outputs=out) ## Wigglegram inp = [gr.Slider(1,15, default = 2, label='StepCycles',step= 1)] with gr.Row(): out = [ gr.Image(type="file", label="Output_images"), #TODO change to gallery gr.Image(type="file", label="Output_wiggle")] btn = gr.Button("Generate Wigglegram") btn.click(fn=wiggle_effect, inputs=inp, outputs=out) demo.launch()