Spaces:
Runtime error
Runtime error
File size: 1,560 Bytes
c7813b3 af4f502 5657662 b32e2df c7813b3 7079dcb af4f502 c7813b3 7079dcb 5657662 2e36fe0 c7813b3 2e36fe0 5657662 c7813b3 7079dcb 6483e22 7079dcb 2e36fe0 7079dcb c7813b3 5657662 bf31859 5657662 7079dcb b32e2df |
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 |
import gradio as gr
import inpaint.infer_model as inpaint
import sod.infer_model as sod
import numpy as np
import torch
import glob
import cv2
# import os
# cmd = 'sh download.sh'
# os.system(cmd)
device = torch.device("cpu")
print(device)
inpaint_model = inpaint.IVModel(device=device)
sod_model = sod.IVModel(device=device)
max_size=512
scale_factor = 8
count = 0
def sod_inpaint(img):
global count
h,w = img.shape[:2]
if max(h, w) > max_size:
if h < w:
h, w = int(max_size * h / w), max_size
else:
h, w = max_size, int(max_size * w / h)
h = h // scale_factor * scale_factor
w = w // scale_factor * scale_factor
img = cv2.resize(img, (w,h))
img = img[:,:,::-1]
sod_res = sod_model.forward(img,None)
sod_res = np.uint8(sod_res)
h,w = sod_res.shape[:2]
so = np.uint8(sod_res[:,:w//2,:] * (sod_res[:,w//2:,:]>0).astype(np.float32))
inpaint_res = inpaint_model.forward(sod_res,None)
inpaint_res = np.uint8(inpaint_res)
count +=1
print(count, ' images have been processed')
return so[:,:,::-1], inpaint_res[:,:,::-1]
examples = glob.glob('examples/*.*')
inputs = gr.inputs.Image(shape=None, image_mode="RGB", invert_colors=False, source="upload", tool="editor", type="numpy", label=None, optional=False)
iface = gr.Interface(fn=sod_inpaint, inputs=inputs, outputs=["image", "image"], examples=examples, title='Salient Object Detection + Inpaint', description='Upload an image and you will see the fg and inpainted bg', theme='huggingface')
iface.launch()
|