Spaces:
Runtime error
Runtime error
File size: 4,751 Bytes
f19da68 ffead1e f19da68 4cae45e 559b00c 4cae45e 559b00c 4cae45e ffead1e 4cae45e |
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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
import gradio as gr
import json
import torch
import time
import random
try:
# Only on HuggingFace
import spaces
is_space_imported = True
except ImportError:
is_space_imported = False
from tqdm import tqdm
from huggingface_hub import snapshot_download
from models import AudioDiffusion, DDPMScheduler
from audioldm.audio.stft import TacotronSTFT
from audioldm.variational_autoencoder import AutoencoderKL
from pydub import AudioSegment
# Old code
import numpy as np
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
from huggingface_hub import hf_hub_download
from gradio_imageslider import ImageSlider
from briarmbg import BriaRMBG
import PIL
from PIL import Image
from typing import Tuple
net=BriaRMBG()
# model_path = "./model1.pth"
#model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
model_path = hf_hub_download("cocktailpeanut/gbmr", 'model.pth')
if torch.cuda.is_available():
net.load_state_dict(torch.load(model_path))
net=net.cuda()
device = "cuda"
elif torch.backends.mps.is_available():
net.load_state_dict(torch.load(model_path,map_location="mps"))
net=net.to("mps")
device = "mps"
else:
net.load_state_dict(torch.load(model_path,map_location="cpu"))
device = "cpu"
net.eval()
def resize_image(image):
image = image.convert('RGB')
model_input_size = (1024, 1024)
image = image.resize(model_input_size, Image.BILINEAR)
return image
def process(image):
# prepare input
orig_image = Image.fromarray(image)
w,h = orig_im_size = orig_image.size
image = resize_image(orig_image)
im_np = np.array(image)
im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
im_tensor = torch.unsqueeze(im_tensor,0)
im_tensor = torch.divide(im_tensor,255.0)
im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
if device == "cuda":
im_tensor=im_tensor.cuda()
elif device == "mps":
im_tensor=im_tensor.to("mps")
#inference
result=net(im_tensor)
# post process
result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
ma = torch.max(result)
mi = torch.min(result)
result = (result-mi)/(ma-mi)
# image to pil
im_array = (result*255).cpu().data.numpy().astype(np.uint8)
pil_im = Image.fromarray(np.squeeze(im_array))
# paste the mask on the original image
new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
new_im.paste(orig_image, mask=pil_im)
# new_orig_image = orig_image.convert('RGBA')
return new_im
# return [new_orig_image, new_im]
# block = gr.Blocks().queue()
# with block:
# gr.Markdown("## BRIA RMBG 1.4")
# gr.HTML('''
# <p style="margin-bottom: 10px; font-size: 94%">
# This is a demo for BRIA RMBG 1.4 that using
# <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">BRIA RMBG-1.4 image matting model</a> as backbone.
# </p>
# ''')
# with gr.Row():
# with gr.Column():
# input_image = gr.Image(sources=None, type="pil") # None for upload, ctrl+v and webcam
# # input_image = gr.Image(sources=None, type="numpy") # None for upload, ctrl+v and webcam
# run_button = gr.Button(value="Run")
# with gr.Column():
# result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", columns=[1], height='auto')
# ips = [input_image]
# run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
# block.launch(debug = True)
# block = gr.Blocks().queue()
gr.Markdown("## BRIA RMBG 1.4")
gr.HTML('''
<p style="margin-bottom: 10px; font-size: 94%">
This is a demo for BRIA RMBG 1.4 that using
<a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">BRIA RMBG-1.4 image matting model</a> as backbone.
</p>
''')
title = "Background Removal"
description = r"""Background removal model developed by <a href='https://BRIA.AI' target='_blank'><b>BRIA.AI</b></a>, trained on a carefully selected dataset and is available as an open-source model for non-commercial use.<br>
For test upload your image and wait. Read more at model card <a href='https://huggingface.co/briaai/RMBG-1.4' target='_blank'><b>briaai/RMBG-1.4</b></a>.<br>
"""
examples = [['./input.jpg'],]
# output = ImageSlider(position=0.5,label='Image without background', type="pil", show_download_button=True)
# demo = gr.Interface(fn=process,inputs="image", outputs=output, examples=examples, title=title, description=description)
demo = gr.Interface(fn=process,inputs="image", outputs="image", examples=examples, title=title, description=description)
if __name__ == "__main__":
demo.launch(share=False)
|