Spaces:
Runtime error
Runtime error
File size: 4,652 Bytes
b145c13 e0cfc5c b145c13 9241860 e0cfc5c 939dab3 e0cfc5c b145c13 9241860 e0cfc5c 939dab3 e0cfc5c b145c13 e0cfc5c b145c13 e0cfc5c b145c13 9241860 e0cfc5c 9241860 b145c13 9241860 b145c13 9241860 b145c13 9241860 b145c13 9241860 b145c13 |
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 135 |
import cv2
import gradio as gr
import os
import requests
from PIL import Image
import numpy as np
import torch
from torch.autograd import Variable
from torchvision import transforms
import torch.nn.functional as F
# Automatically download required files
# 1. data_loader_cache.py from GitHub
if not os.path.exists("data_loader_cache.py"):
print("Downloading data_loader_cache.py...")
try:
response = requests.get("https://raw.githubusercontent.com/xuebinqin/DIS/main/DIS/IS-Net/data_loader_cache.py")
response.raise_for_status()
with open("data_loader_cache.py", "wb") as f:
f.write(response.content)
except requests.RequestException as e:
print(f"Failed to download data_loader_cache.py: {e}")
raise
# 2. models.py from GitHub
if not os.path.exists("models.py"):
print("Downloading models.py...")
try:
response = requests.get("https://raw.githubusercontent.com/xuebinqin/DIS/main/DIS/IS-Net/models.py")
response.raise_for_status()
with open("models.py", "wb") as f:
f.write(response.content)
except requests.RequestException as e:
print(f"Failed to download models.py: {e}")
raise
# 3. isnet.pth from Hugging Face Git LFS (direct URL from screenshot)
if not os.path.exists("saved_models"):
os.makedirs("saved_models")
isnet_path = "saved_models/isnet.pth"
if not os.path.exists(isnet_path):
print("Downloading isnet.pth from Hugging Face Git LFS...")
try:
lfs_url = "https://cdn-lfs.huggingface.co/repos/e0/a8/e0a889743a78391b48db7c4c0b4de1963ee320cb10934c75a32481dc5af9c61/e0a889743a78391b48db7c4c0b4de1963ee320cb10934c75a32481dc5af9c61?download=true"
response = requests.get(lfs_url, stream=True)
response.raise_for_status()
with open(isnet_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
except requests.RequestException as e:
print(f"Failed to download isnet.pth: {e}")
raise
# Project imports
from data_loader_cache import normalize, im_reader, im_preprocess
from models import *
# Helpers
device = 'cpu'
class GOSNormalize(object):
def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
self.mean = mean
self.std = std
def __call__(self, image):
image = normalize(image, self.mean, self.std)
return image
transform = transforms.Compose([GOSNormalize([0.5, 0.5, 0.5], [1.0, 1.0, 1.0])])
def load_image(im_path, hypar):
im = im_reader(im_path)
im, im_shp = im_preprocess(im, hypar["cache_size"])
im = torch.divide(im, 255.0)
shape = torch.from_numpy(np.array(im_shp))
return transform(im).unsqueeze(0), shape.unsqueeze(0)
def build_model(hypar, device):
net = hypar["model"]
net.to(device)
if hypar["restore_model"]:
net.load_state_dict(torch.load(os.path.join(hypar["model_path"], hypar["restore_model"]), map_location=device))
net.eval()
return net
def predict(net, inputs_val, shapes_val, hypar, device):
net.eval()
inputs_val = inputs_val.type(torch.FloatTensor).to(device)
with torch.no_grad():
inputs_val_v = Variable(inputs_val)
ds_val = net(inputs_val_v)[0]
pred_val = ds_val[0][0, :, :, :]
pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val, 0), (shapes_val[0][0], shapes_val[0][1]), mode='bilinear'))
ma = torch.max(pred_val)
mi = torch.min(pred_val)
pred_val = (pred_val - mi) / (ma - mi)
return (pred_val.cpu().numpy() * 255).astype(np.uint8)
# Set Parameters
hypar = {
"model_path": "saved_models",
"restore_model": "isnet.pth",
"cache_size": [512, 512],
"input_size": [512, 512],
"crop_size": [512, 512],
"model": ISNetDIS()
}
# Build Model
net = build_model(hypar, device)
def inference(image):
image_path = image
image_tensor, orig_size = load_image(image_path, hypar)
mask = predict(net, image_tensor, orig_size, hypar, device)
pil_mask = Image.fromarray(mask).convert('L')
im_rgb = Image.open(image).convert("RGB")
im_rgba = im_rgb.copy()
im_rgba.putalpha(pil_mask)
return [im_rgba, pil_mask]
title = "Dichotomous Image Segmentation"
description = "Upload an image to remove its background."
interface = gr.Interface(
fn=inference,
inputs=gr.Image(type='filepath'),
outputs=[gr.Image(type='filepath', format="png"), gr.Image(type='filepath', format="png")],
title=title,
description=description,
flagging_mode="never",
cache_mode="lazy"
).launch() |