segmentngo / app.py
dputilov's picture
Update app.py
7e18798 verified
try:
import detectron2
except:
import os
os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
from git import Repo
from PIL import Image
import os
import logging
import io
import cv2
from io import BytesIO
import gradio as gr
import numpy as np
PAT_TOKEN = os.environ.get("PAT_TOKEN")
GIT_REPO = os.environ.get("GIT_REPO")
if os.path.exists('./segmentspace'):
pass
else:
Repo.clone_from(f"https://{PAT_TOKEN}:@github.com/{GIT_REPO}",
"./segmentspace")
from segmentspace import Remover
MAX_FILE_SIZE = 5 * 1024 * 1024
remover = Remover()
def inference(image):
image_path = image
image = Image.open(image_path).convert('RGB')
output_image, output_mask = remover.process(image)
pil_mask = Image.fromarray(output_mask).convert('L')
im_rgb = Image.fromarray(output_image).convert("RGB")
im_rgba = im_rgb.copy()
im_rgba.putalpha(pil_mask)
return im_rgba
title = "Segment 'n Go"
description = "To use it, simply upload your image, or click one of the examples to load them."
interface = gr.Interface(
fn=inference,
inputs=[gr.Image(type='filepath')],
outputs=[gr.Image(type='pil')],
examples=[['zebra.jpg'], ['zebra.jpg']],
title=title,
description=description,
allow_flagging='never',
cache_examples=False,
).queue().launch(show_error=True)