app.py

import gradio as gr
from rembg import remove
from PIL import Image
import numpy as np
import io

# Hàm xử lý tách nền
def process_rembg(input_img, bg_setting, progress=gr.Progress()):
    if input_img is None:
        return None
    
    progress(0.2, desc="🔍 Đang quét chủ thể...")
    
    # Chuyển sang mảng numpy để rembg xử lý
    img_array = np.array(input_img)
    
    # Thực hiện tách nền (RemBg tự động dùng CPU)
    output_array = remove(img_array)
    no_bg_img = Image.fromarray(output_array)
    
    progress(0.6, desc="🎨 Đang áp dụng màu nền...")
    
    # Xử lý các tùy chọn nền
    if bg_setting == "Trong suốt (PNG)":
        progress(1.0, desc="✅ Xong!")
        return no_bg_img
    
    # Bản đồ màu sắc
    colors = {
        "Trắng (White)": (255, 255, 255),
        "Đen (Black)": (0, 0, 0),
        "Xanh thẻ (Blue)": (0, 102, 204),
        "Xanh lá (Green)": (0, 204, 102),
        "Đỏ (Red)": (204, 0, 0)
    }
    
    # Tạo nền mới và dán ảnh đã tách lên
    new_bg = Image.new("RGB", no_bg_img.size, colors.get(bg_setting, (255, 255, 255)))
    new_bg.paste(no_bg_img, (0, 0), no_bg_img)
    
    progress(1.0, desc="✅ Hoàn thành!")
    return new_bg

# Giao diện Blocks cho chuyên nghiệp
css = """
footer {visibility: hidden}
.gradio-container {max-width: 900px !important; margin: auto}
#header {text-align: center; margin-bottom: 20px}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.HTML("<div id='header'><h1>✂️ AI Background Remover Pro</h1><p>Tách nền ảnh nhanh - Chuẩn - Miễn phí</p></div>")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Tải ảnh của bạn lên")
            bg_type = gr.Dropdown(
                choices=["Trong suốt (PNG)", "Trắng (White)", "Đen (Black)", "Xanh thẻ (Blue)", "Xanh lá (Green)", "Đỏ (Red)"],
                value="Trong suốt (PNG)",
                label="Chọn loại nền kết quả"
            )
            submit_btn = gr.Button("🚀 Bắt đầu tách nền", variant="primary")
            
        with gr.Column(scale=1):
            output_image = gr.Image(type="pil", label="Kết quả xử lý")
            gr.Markdown("📝 **Lưu ý:** Lần đầu chạy sẽ mất khoảng 30s để AI tải model về hệ thống.")

    submit_btn.click(
        fn=process_rembg, 
        inputs=[input_image, bg_type], 
        outputs=output_image
    )

if __name__ == "__main__":
    demo.launch()