Rename main.py to app.py

app.py

@@ -0,0 +1,97 @@
+import os
+import cv2
+import numpy as np
+import onnxruntime as ort
+import uuid
+import base64
+from PIL import Image
+import gradio as gr
+import spaces
+
+# Setup
+API_KEY = os.getenv("API_KEY", "demo")
+INPUT_SIZE = (512, 512)
+MODEL_PATH = "BiRefNet-general-resolution_512x512-fp16-epoch_216.onnx"
+
+# Load ONNX model
+assert os.path.exists(MODEL_PATH), f"Model not found: {MODEL_PATH}"
+session = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
+input_name = session.get_inputs()[0].name
+
+# Preprocess
+def preprocess_image(image: np.ndarray):
+    original_shape = image.shape[:2]
+    resized = cv2.resize(image, INPUT_SIZE)
+    normalized = (resized.astype(np.float32) / 255.0 - 0.5) / 0.5
+    transposed = np.transpose(normalized, (2, 0, 1))
+    input_tensor = np.expand_dims(transposed, axis=0).astype(np.float32)
+    return input_tensor, original_shape, image
+
+# Mask logic
+def apply_mask(original_img, mask_array, original_shape):
+    mask = np.squeeze(mask_array)
+    resized_mask = cv2.resize(mask, (original_shape[1], original_shape[0]))
+    binary_mask = (resized_mask > 0.5).astype(np.uint8)
+    alpha = (binary_mask * 255).astype(np.uint8)
+    masked = cv2.bitwise_and(original_img, original_img, mask=binary_mask)
+    bgra = cv2.cvtColor(masked, cv2.COLOR_RGB2BGRA)
+    bgra[:, :, 3] = alpha
+    return bgra
+
+# ============ UI ============
+@spaces.GPU
+def remove_background_ui(image, bg=None):
+    input_tensor, original_shape, original_img = preprocess_image(image)
+    mask = session.run(None, {input_name: input_tensor})[0]
+    result = apply_mask(original_img, mask, original_shape)
+
+    if bg is not None:
+        bg_resized = cv2.resize(bg, (original_shape[1], original_shape[0]))
+        alpha = result[:, :, 3] / 255.0
+        fg = result[:, :, :3]
+        blended = (fg * alpha[..., None] + bg_resized * (1 - alpha[..., None])).astype(np.uint8)
+        return Image.fromarray(blended)
+    return Image.fromarray(result)
+
+# ============ API ============
+@spaces.GPU
+def remove_background_api(image_file, api_key=""):
+    if api_key != API_KEY:
+        raise gr.Error("❌ Invalid API Key")
+    image = Image.open(image_file).convert("RGB")
+    image_np = np.array(image)
+    input_tensor, original_shape, original_img = preprocess_image(image_np)
+    mask = session.run(None, {input_name: input_tensor})[0]
+    result = apply_mask(original_img, mask, original_shape)
+    success, buffer = cv2.imencode(".png", result)
+    return f"data:image/png;base64,{base64.b64encode(buffer).decode('utf-8')}"
+
+# Gradio interfaces
+ui = gr.Interface(
+    fn=remove_background_ui,
+    inputs=[
+        gr.Image(type="numpy", label="Main Image"),
+        gr.Image(type="numpy", label="Optional Background")
+    ],
+    outputs=gr.Image(type="pil", label="Result"),
+    title="🖼️ Background Remover",
+    description="Upload a photo (and optionally a background)."
+)
+
+api = gr.Interface(
+    fn=remove_background_api,
+    inputs=[
+        gr.Image(type="filepath", label="Upload Image"),
+        gr.Text(label="API Key", type="password")
+    ],
+    outputs=gr.Text(label="Base64 PNG"),
+    title="🔐 API Access",
+    description="POST to `/run/predict` with file + API key."
+)
+
+# Final Gradio app with predict support
+demo = gr.TabbedInterface([ui, api], ["Web UI", "API Access"])
+
+# Launch
+if __name__ == "__main__":
+    demo.launch()

main.py

@@ -1,248 +0,0 @@
-###############################################################################
-# 1. Set environment variables BEFORE importing Gradio
-###############################################################################
-import os
-
-os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
-os.environ["GRADIO_SERVER_PORT"] = "7860"
-os.environ["GRADIO_ROOT_PATH"] = "/"
-
-###############################################################################
-# 2. Imports
-###############################################################################
-import uuid
-import base64
-import json
-import shutil
-import traceback
-import numpy as np
-import cv2
-import onnxruntime as ort
-from io import BytesIO
-from PIL import Image
-from datetime import datetime
-from pathlib import Path
-
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Depends, Request
-from fastapi.responses import FileResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-from fastapi.middleware.cors import CORSMiddleware
-
-import gradio as gr
-import spaces
-import uvicorn
-
-###############################################################################
-# 3. Setup
-###############################################################################
-API_KEY = os.getenv("API_KEY")
-
-app = FastAPI(title="Background Removal API")
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-TMP_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")
-os.makedirs(TMP_FOLDER, exist_ok=True)
-print(f"Created tmp folder at: {TMP_FOLDER}")
-
-app.mount("/tmp", StaticFiles(directory=TMP_FOLDER), name="tmp")
-templates = Jinja2Templates(directory="templates")
-
-model_path = "BiRefNet-general-resolution_512x512-fp16-epoch_216.onnx"
-session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
-assert "CUDAExecutionProvider" not in ort.get_available_providers(), \
-    "CUDA provider found but not supported on ZeroGPU."
-
-input_name = session.get_inputs()[0].name
-INPUT_SIZE = (512, 512)
-
-###############################################################################
-# 4. Utilities
-###############################################################################
-
-def verify_api_key(api_key: str = Form(...)):
-    if api_key != API_KEY:
-        raise HTTPException(status_code=401, detail="Invalid API key")
-    return api_key
-@spaces.GPU
-def preprocess_image(image):
-    if isinstance(image, str):
-        img = cv2.imread(image)
-    elif isinstance(image, np.ndarray):
-        img = image
-    else:
-        nparr = np.frombuffer(image, np.uint8)
-        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-
-    original_img = img.copy()
-    original_shape = img.shape[:2]
-    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    resized = cv2.resize(rgb, INPUT_SIZE)
-    normalized = resized.astype(np.float32) / 255.0
-    normalized = (normalized - 0.5) / 0.5
-    transposed = np.transpose(normalized, (2, 0, 1))
-    input_tensor = np.expand_dims(transposed, axis=0).astype(np.float32)
-
-    return input_tensor, original_shape, original_img
-
-@spaces.GPU(duration=240)
-def apply_mask(original_img, mask_array, original_shape, output_path):
-    try:
-        mask = np.squeeze(mask_array)
-        mask = cv2.resize(mask, (original_shape[1], original_shape[0]))
-        mask = np.clip(mask, 0, 1)
-        binary_mask = (mask > 0.5).astype(np.uint8)
-
-        img = original_img.astype(np.uint8)
-        masked_img = cv2.bitwise_and(img, img, mask=binary_mask)
-        alpha = (binary_mask * 255).astype(np.uint8)
-
-        bgra = cv2.cvtColor(masked_img, cv2.COLOR_BGR2BGRA)
-        bgra[:, :, 3] = alpha
-
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
-        cv2.imwrite(output_path, bgra, [cv2.IMWRITE_PNG_COMPRESSION, 0])
-        return bgra, True
-    except Exception as e:
-        print(f"Error applying mask: {e}")
-        return None, False
-
-###############################################################################
-# 5. FastAPI Routes
-###############################################################################
-@app.post("/")
-async def index_post(request: Request, main_photo: UploadFile = File(...), bg_photo: UploadFile = File(None)):
-    try:
-        main_image_data = await main_photo.read()
-        input_tensor, original_shape, original_img = preprocess_image(main_image_data)
-        output = session.run(None, {input_name: input_tensor})
-        mask = output[0]
-
-        result_filename = f"{uuid.uuid4()}.png"
-        output_path = os.path.join(TMP_FOLDER, result_filename)
-        transparent_img, success = apply_mask(original_img, mask, original_shape, output_path)
-        final_result_path = output_path
-
-        if bg_photo:
-            bg_image_data = await bg_photo.read()
-            bg_np = np.frombuffer(bg_image_data, np.uint8)
-            bg_img = cv2.imdecode(bg_np, cv2.IMREAD_COLOR)
-            bg_img_resized = cv2.resize(bg_img, (original_shape[1], original_shape[0]))
-            alpha = transparent_img[:, :, 3] / 255.0
-            foreground = transparent_img[:, :, :3]
-            blended = (foreground * alpha[..., None] + bg_img_resized * (1 - alpha[..., None])).astype(np.uint8)
-            final_result_path = os.path.join(TMP_FOLDER, f"bg_replaced_{uuid.uuid4()}.png")
-            cv2.imwrite(final_result_path, blended)
-
-        return templates.TemplateResponse("index.html", {"request": request, "output_image": os.path.basename(final_result_path)})
-    except Exception as e:
-        print(traceback.format_exc())
-        return templates.TemplateResponse("index.html", {"request": request, "error": f"Error: {str(e)}"})
-
-@app.get("/")
-async def index_get(request: Request):
-    return templates.TemplateResponse("index.html", {"request": request})
-
-@app.post("/remove-background")
-async def remove_background(request: Request, api_key: str = Form(...), main_photo: UploadFile = File(...)):
-    verify_api_key(api_key)
-    try:
-        image_data = await main_photo.read()
-        result_filename = f"{uuid.uuid4()}.png"
-        output_path = os.path.join(TMP_FOLDER, result_filename)
-        input_tensor, original_shape, original_img = preprocess_image(image_data)
-        output = session.run(None, {input_name: input_tensor})
-        mask = output[0]
-        _, success = apply_mask(original_img, mask, original_shape, output_path)
-
-        if success:
-            base_url = str(request.base_url).rstrip("/")
-            image_url = f"{base_url}/tmp/{result_filename}"
-            return {"status": "success", "message": "Background removed", "filename": result_filename, "image_url": image_url}
-        return {"status": "failure", "message": "Failed to process image"}
-    except Exception as e:
-        print(traceback.format_exc())
-        return {"status": "failure", "message": f"Error: {str(e)}"}
-
-@app.post("/process_image")
-async def process_image(request: Request, image: UploadFile = File(...), api_key: str = Form(...)):
-    verify_api_key(api_key)
-    try:
-        image_data = await image.read()
-        result_filename = f"{uuid.uuid4()}.png"
-        output_path = os.path.join(TMP_FOLDER, result_filename)
-        input_tensor, original_shape, original_img = preprocess_image(image_data)
-        output = session.run(None, {input_name: input_tensor})
-        mask = output[0]
-        bgra, success = apply_mask(original_img, mask, original_shape, output_path)
-
-        if success:
-            with open(output_path, "rb") as img_file:
-                base64_image = base64.b64encode(img_file.read()).decode("utf-8")
-            return {"status": "success", "image_code": f"data:image/png;base64,{base64_image}"}
-        return {"status": "failure", "message": "Failed to process image"}
-    except Exception as e:
-        print(traceback.format_exc())
-        return {"status": "failure", "message": f"Error: {str(e)}"}
-
-@app.get("/download/{filename}")
-async def download_file(filename: str):
-    file_path = os.path.join(TMP_FOLDER, filename)
-    if os.path.exists(file_path):
-        return FileResponse(path=file_path, filename=filename, media_type="image/png")
-    raise HTTPException(status_code=404, detail="File not found")
-
-###############################################################################
-# 6. Gradio Interface
-###############################################################################
-@spaces.GPU
-def process_image_gradio(image):
-    session = ort.InferenceSession(
-        model_path,
-        providers=["CUDAExecutionProvider"]
-    )
-    input_tensor, original_shape, original_img = preprocess_image(image)
-    output = session.run(None, {input_name: input_tensor})
-    mask = output[0]
-    result_filename = f"{uuid.uuid4()}.png"
-    output_path = os.path.join(TMP_FOLDER, result_filename)
-    result_img, success = apply_mask(original_img, mask, original_shape, output_path)
-    if success:
-        return Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_BGRA2RGBA))
-    return None
-process_image_gradio.zerogpu =True
-# Optional: this ensures startup detects GPU
-@spaces.GPU(duration=5)
-def _gpu_wakeup():
-    return "GPU ready"
-
-# Gradio Blocks Interface
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Background Removal (GPU)")
-    gr.Markdown("Upload an image to remove the background using ONNX on GPU")
-
-    with gr.Row():
-        image_input = gr.Image(type="numpy", label="Upload Image")
-        image_output = gr.Image(type="pil", label="Output")
-
-    submit_btn = gr.Button("Remove Background")
-    submit_btn.click(fn=process_image_gradio, inputs=image_input, outputs=image_output)
-
-
-###############################################################################
-# 7. Mount Gradio App
-###############################################################################
-gr.mount_gradio_app(app, demo, path="/", ssr_mode=False)
-
-###############################################################################
-# 8. Run (only needed for local dev; not needed on Hugging Face)
-###############################################################################
-if __name__ == "__main__":
-   uvicorn.run(app, host="0.0.0.0", port=7860)