|
|
import torch |
|
|
from PIL import Image |
|
|
import numpy as np |
|
|
import gradio as gr |
|
|
from transformers import AutoFeatureExtractor |
|
|
from huggingface_hub import snapshot_download |
|
|
import sys |
|
|
import os |
|
|
|
|
|
|
|
|
model_path = snapshot_download(repo_id="briaai/RMBG-1.4", trust_remote_code=True) |
|
|
|
|
|
|
|
|
sys.path.insert(0, model_path) |
|
|
|
|
|
|
|
|
from model import RMBGModel |
|
|
|
|
|
|
|
|
model = RMBGModel.from_pretrained(model_path, trust_remote_code=True) |
|
|
model.eval() |
|
|
|
|
|
|
|
|
extractor = AutoFeatureExtractor.from_pretrained("briaai/RMBG-1.4") |
|
|
|
|
|
def process(image): |
|
|
inputs = extractor(images=image, return_tensors="pt") |
|
|
with torch.no_grad(): |
|
|
result = model(**inputs) |
|
|
alpha = result["output"].squeeze().numpy() |
|
|
alpha = (alpha * 255).astype(np.uint8) |
|
|
alpha = Image.fromarray(alpha).resize(image.size) |
|
|
|
|
|
image = image.convert("RGBA") |
|
|
alpha = alpha.convert("L") |
|
|
image.putalpha(alpha) |
|
|
return image |
|
|
|
|
|
demo = gr.Interface( |
|
|
fn=process, |
|
|
inputs=gr.Image(type="pil"), |
|
|
outputs=gr.Image(type="pil"), |
|
|
title="إزالة خلفية الصور باستخدام RMBG", |
|
|
description="ارفع صورة وسيتم إزالة الخلفية تلقائيًا باستخدام نموذج RMBG من briaai." |
|
|
) |
|
|
|
|
|
demo.launch() |
|
|
|
