WiggleGAN / app.py
Rodrigo_Cobo
fix jpeg error
eac489d
raw
history blame
1.78 kB
import os
import gradio as gr
import cv2
import torch
import urllib.request
from PIL import Image
import matplotlib.pyplot as plt
def update(slider, img):
if not os.path.exists('temp'):
os.system('mkdir temp')
filename = "temp/image.jpg"
img.save(filename, "JPEG")
model_type = "DPT_Hybrid"
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()
#out_im = Image.fromarray(output)
#out_im.convert('RGB').save("temp/image_depth.jpeg", "JPEG")
cv2.imwrite("temp/image_depth.jpeg", output)
return f'temp/image_depth.jpeg'
with gr.Blocks() as demo:
gr.Markdown("Start typing below and then click **Run** to see the output.")
inp = [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="file", label="Output")
btn = gr.Button("Run")
btn.click(fn=update, inputs=inp, outputs=out)
demo.launch()