Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import cv2 | |
| from rembg import new_session, remove | |
| from rembg.sessions import sessions_class | |
| def inference(file, mask, model, x, y): | |
| im = cv2.imread(file, cv2.IMREAD_COLOR) | |
| input_path = "input.png" | |
| output_path = "output.png" | |
| cv2.imwrite(input_path, im) | |
| with open(input_path, 'rb') as i: | |
| with open(output_path, 'wb') as o: | |
| input = i.read() | |
| session = new_session(model) | |
| output = remove( | |
| input, | |
| session=session, | |
| **{ "sam_prompt": [{"type": "point", "data": [x, y], "label": 1}] }, | |
| only_mask=(mask == "Mask only") | |
| ) | |
| o.write(output) | |
| return output_path | |
| title = "RemBG" | |
| description = "Gradio demo for **[RemBG](https://github.com/danielgatis/rembg)**. To use it, simply upload your image, select a model, click Process, and wait." | |
| badge = """ | |
| <div style="position: fixed; left: 50%; text-align: center;"> | |
| <a href="https://github.com/danielgatis/rembg" target="_blank" style="text-decoration: none;"> | |
| <img src="https://img.shields.io/badge/RemBG-Github-blue" alt="RemBG Github" /> | |
| </a> | |
| </div> | |
| """ | |
| def get_coords(evt: gr.SelectData) -> tuple: | |
| return evt.index[0], evt.index[1] | |
| def show_coords(model: str): | |
| visible = model == "sam" | |
| return gr.update(visible=visible), gr.update(visible=visible), gr.update(visible=visible) | |
| for session in sessions_class: | |
| session.download_models() | |
| with gr.Blocks() as app: | |
| gr.Markdown(f"# {title}") | |
| gr.Markdown(description) | |
| with gr.Row(): | |
| inputs = gr.Image(type="filepath", label="Input Image") | |
| outputs = gr.Image(type="filepath", label="Output Image") | |
| with gr.Row(): | |
| mask_option = gr.Radio( | |
| ["Default", "Mask only"], | |
| value="Default", | |
| label="Output Type" | |
| ) | |
| model_selector = gr.Dropdown( | |
| [ | |
| "u2net", | |
| "u2netp", | |
| "u2net_human_seg", | |
| "u2net_cloth_seg", | |
| "silueta", | |
| "isnet-general-use", | |
| "isnet-anime", | |
| "sam", | |
| "birefnet-general", | |
| "birefnet-general-lite", | |
| "birefnet-portrait", | |
| "birefnet-dis", | |
| "birefnet-hrsod", | |
| "birefnet-cod", | |
| "birefnet-massive" | |
| ], | |
| value="isnet-general-use", | |
| label="Model Selection" | |
| ) | |
| extra = gr.Markdown("## Click on the image to capture coordinates (for SAM model)", visible=False) | |
| x = gr.Number(label="Mouse X Coordinate", visible=False) | |
| y = gr.Number(label="Mouse Y Coordinate", visible=False) | |
| model_selector.change(show_coords, inputs=model_selector, outputs=[x, y, extra]) | |
| inputs.select(get_coords, None, [x, y]) | |
| gr.Button("Process Image").click( | |
| inference, | |
| inputs=[inputs, mask_option, model_selector, x, y], | |
| outputs=outputs | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["lion.png", "Default", "u2net", None, None], | |
| ["girl.jpg", "Default", "u2net", None, None], | |
| ["anime-girl.jpg", "Default", "isnet-anime", None, None] | |
| ], | |
| inputs=[inputs, mask_option, model_selector, x, y], | |
| outputs=outputs | |
| ) | |
| gr.HTML(badge) | |
| app.launch() | |