Update app.py

app.py

@@ -1,9 +1,15 @@
+import numpy as np
 import torch
 import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
 from huggingface_hub import hf_hub_download
 from briarmbg import BriaRMBG
 from PIL import Image
-import gradio as gr
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import FileResponse
+import os
+
+app = FastAPI()
 
 # 모델 로드
 net = BriaRMBG()
@@ -21,8 +27,8 @@ def resize_image(image):
     image = image.resize(model_input_size, Image.BILINEAR)
     return image
 
-def process(image):
-    orig_image = Image.fromarray(image)
+def process_image(image: Image.Image):
+    orig_image = image
     w, h = orig_im_size = orig_image.size
     image = resize_image(orig_image)
     im_np = np.array(image)
@@ -50,17 +56,14 @@ def process(image):
     output_path = "output_image.png"
     new_im.save(output_path)
     
-    return new_im, output_path
+    return output_path
 
-# Gradio 인터페이스 설정
-demo = gr.Interface(
-    fn=process, 
-    inputs="image", 
-    outputs=["image", "file"], 
-    title="이미지 배경 제거", 
-    description="이미지의 배경을 제거하고 결과 이미지를 다운로드할 수 있는 데모입니다."
-)
+@app.post("/remove-background/")
+async def remove_background(file: UploadFile = File(...)):
+    image = Image.open(file.file)
+    output_path = process_image(image)
+    return FileResponse(output_path, media_type="image/png", filename="output_image.png")
 
 if __name__ == "__main__":
-    # Gradio 서버 실행
-    demo.queue(max_size=12).launch(server_name="0.0.0.0", server_port=7910, show_api=True)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7910)