learn-2

Running

App Files Files Community

urnotwen commited on 20 days ago

Commit

a2e0537

verified ·

1 Parent(s): d9af6a3

Update app.py

Browse files

Files changed (1) hide show

app.py +46 -61

app.py CHANGED Viewed

@@ -1,99 +1,84 @@
 import gradio as gr
 from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
 from PIL import Image
-import io
-# --- 1. 初始化模型 (只會在啟動時執行一次) ---
 model_id = "briaai/RMBG-2.0"
 print(f"正在載入模型: {model_id} ...")
 try:
-    # 載入模型到 CPU (Hugging Face Free Tier 只有 CPU)
-    model = AutoModelForImageSegmentation.from_pretrained(model_id, trust_remote_code=True)
-    device = torch.device("cpu") # 強制使用 CPU
     model.to(device)
-    model.eval() # 設定為評估模式
-    print("模型載入成功！")
 except Exception as e:
-    print(f"模型載入失敗: {e}")
-# --- 2. 定義圖像處理邏輯 ---
 def process_image(input_image):
     if input_image is None:
         return None
-    # 紀錄原始尺寸
-    orig_w, orig_h = input_image.size
-    # 準備輸入 (RMBG 2.0 建議尺寸為 1024x1024)
     transform_image = transforms.Compose([
-        transforms.Resize((1024, 1024)),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])
-    input_tensor = transform_image(input_image).unsqueeze(0).to(device)
-    # 推理 (Inference)
     with torch.no_grad():
-        preds = model(input_tensor)[0][0]
-        # 還原到原始尺寸
-        preds = torch.nn.functional.interpolate(preds, size=(orig_h, orig_w), mode='bilinear', align_corners=False)
-        preds = torch.sigmoid(preds)
-    # 處理遮罩
-    mask = preds.squeeze().cpu().numpy()
-    # 轉回 PIL Image
-    mask_img = Image.fromarray((mask * 255).astype('uint8'), mode='L')
     # 合成去背圖
-    output_img = input_image.convert("RGBA")
-    output_img.putalpha(mask_img)
-    return output_img
-# --- 3. 設定 PWA 與手機優化 HTML ---
-# 這些標籤會讓網頁在「加入主畫面」後變成全螢幕 APP
-pwa_header = """
-<head>
-    <meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no">
-    <meta name="apple-mobile-web-app-capable" content="yes">
-    <meta name="apple-mobile-web-app-status-bar-style" content="black-translucent">
-    <meta name="theme-color" content="#0b0f19">
-    <title>AI 去背神器</title>
-    <style>
-        /* 隱藏 Gradio 預設的頁尾，讓畫面更乾淨 */
-        footer {display: none !important;}
-        .gradio-container {min-height: 100vh !important;}
-    </style>
-</head>
-"""
-# --- 4. 建立 Gradio 介面 ---
-with gr.Blocks(head=pwa_header, theme=gr.themes.Soft()) as app:
-    gr.Markdown(
-        """
-        # ✂️ AI 自動去背 (RMBG 2.0)
-        上傳照片，自動去除背景。
-        """
-    )
-    with gr.Row():
-        # 手機上通常是直向排列，Gradio 會自動響應
-        with gr.Column():
-            input_img = gr.Image(type="pil", label="點擊上傳或拍照", sources=["upload", "clipboard"])
-            btn = gr.Button("開始去背", variant="primary", size="lg")
-        with gr.Column():
-            output_img = gr.Image(type="pil", label="去背結果 (長按儲存)", format="png", show_download_button=True)
-    # 按鈕事件
     btn.click(fn=process_image, inputs=input_img, outputs=output_img)
-# 啟動應用
 if __name__ == "__main__":
     app.launch()

+import os
 import gradio as gr
 from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
 from PIL import Image
+# --- 1. 初始化模型 ---
 model_id = "briaai/RMBG-2.0"
 print(f"正在載入模型: {model_id} ...")
+hf_token = os.getenv("HF_TOKEN")
+# 檢查 Token
+if not hf_token:
+    print("⚠️ 警告: 未偵測到 HF_TOKEN，請檢查 Settings 中的 Secret。")
 try:
+    # 載入模型
+    model = AutoModelForImageSegmentation.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+        token=hf_token
+    )
+    # 強制使用 CPU (避免免費空間報錯)
+    device = torch.device("cpu")
     model.to(device)
+    model.eval()
+    print("✅ 模型載入成功！")
 except Exception as e:
+    print(f"❌ 模型載入失敗: {e}")
+# --- 2. 定義圖像處理邏輯 (官方邏輯) ---
 def process_image(input_image):
     if input_image is None:
         return None
+    # 準備圖像尺寸
+    image_size = (1024, 1024)
+    # 定義轉換
     transform_image = transforms.Compose([
+        transforms.Resize(image_size),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])
+    # 轉換並增加維度
+    input_images = transform_image(input_image).unsqueeze(0).to(device)
+    # 開始預測
     with torch.no_grad():
+        preds = model(input_images)[-1].sigmoid().cpu()
+    # 處理預測結果
+    pred = preds[0].squeeze()
+    # 轉回 PIL 圖片
+    pred_pil = transforms.ToPILImage()(pred)
+    # 調整回原始圖片的大小
+    mask = pred_pil.resize(input_image.size)
     # 合成去背圖
+    image = input_image.convert("RGBA")
+    image.putalpha(mask)
+    return image
+# --- 3. 建立介面 (移除導致錯誤的 head 參數) ---
+with gr.Blocks(theme=gr.themes.Soft()) as app:
+    gr.Markdown("## ✂️ AI 自動去背 (RMBG 2.0)")
+    gr.Markdown("由於版本相容性，目前以標準網頁模式運行。")
+    with gr.Column():
+        input_img = gr.Image(type="pil", label="上傳圖片", sources=["upload", "clipboard", "webcam"])
+        btn = gr.Button("✨ 開始去背", variant="primary", size="lg")
+        output_img = gr.Image(type="pil", label="去背結果", format="png", show_download_button=True)
     btn.click(fn=process_image, inputs=input_img, outputs=output_img)
 if __name__ == "__main__":
     app.launch()