Update app.py

app.py

@@ -1,34 +1,49 @@
 import gradio as gr
-from transformers import pipeline
 import numpy as np
 from PIL import Image
+import onnxruntime as ort
 import requests
 from io import BytesIO
 
-# Ładuj model (automatycznie CPU)
-pipe = pipeline("image-segmentation", 
-                model="schirrmacher/ormbg",
-                trust_remote_code=True)
+# Ładuj model ONNX bezpośrednio (CPU)
+MODEL_ID = "schirrmacher/ormbg"
+ort_session = ort.InferenceSession(f"https://huggingface.co/{MODEL_ID}/resolve/main/model.onnx")
+
+def preprocess_image(image):
+    if isinstance(image, str):  # URL
+        image = Image.open(requests.get(image, stream=True).raw)
+    image = image.convert("RGB").resize((1024, 1024))
+    img_array = np.array(image).astype(np.float32) / 255.0
+    img_array = np.transpose(img_array, (2, 0, 1))  # HWC -> CHW
+    img_array = np.expand_dims(img_array, 0)  # Batch
+    return img_array
 
 def remove_bg(image):
     if image is None:
         return None
     
-    # Usuń tło
-    mask = pipe(image)[0]
-    mask = Image.fromarray(mask['mask'].astype(np.uint8) * 255)
+    # Preprocess
+    input_tensor = preprocess_image(image)
+    
+    # Inferencja ONNX
+    outputs = ort_session.run(None, {"input": input_tensor})
+    mask = outputs[0][0, 0]  # [H, W]
+    
+    # Postprocess
+    mask = (mask > 0.5).astype(np.uint8) * 255
+    mask_pil = Image.fromarray(mask).resize(image.size)
     
     # Zastosuj maskę
-    result = Image.fromarray(np.array(image) * np.array(mask)[:,:,None])
-    return result
+    result = np.array(image) * np.stack([mask_pil, mask_pil, mask_pil], axis=2)
+    return Image.fromarray(result.astype(np.uint8))
 
-# UI Gradio
-iface = gr.Interface(
+# Gradio UI
+demo = gr.Interface(
     fn=remove_bg,
-    inputs=gr.Image(type="pil", label="Upload zdjęcie"),
-    outputs=gr.Image(type="pil", label="Bez tła (PNG)"),
-    title="🖼️ ORMBG - Szybki usuwacz tła (CPU)",
-    description="Model schirrmacher/ormbg - ~400ms na CPU, precyzyjne krawędzie",
+    inputs=gr.Image(type="pil", label="📁 Upload zdjęcie"),
+    outputs=gr.Image(type="pil", label="✅ Bez tła (PNG)"),
+    title="⚡ ORMBG - Szybki CPU Background Remover",
+    description="schirrmacher/ormbg (~350ms CPU) - precyzyjne krawędzie, włos/futro",
     examples=[
         "https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d?w=512",
         "https://images.unsplash.com/photo-1544725176-7c40e5a71c5e?w=512"
@@ -36,4 +51,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()