Update app.py

app.py

@@ -1,37 +1,16 @@
 import gradio as gr
-from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
-from PIL import Image
-import torch
-
-# 모델과 이미지 프로세서 로드
-model_name = "nvidia/mit-b0"
-image_processor = SegformerImageProcessor.from_pretrained(model_name)
-model = SegformerForSemanticSegmentation.from_pretrained(model_name)
+from rembg import remove
 
 def remove_background(image):
-    # 이미지에서 피처 추출
-    inputs = image_processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
-    # 결과에서 가장 가능성 높은 클래스 선택
-    upsampled_logits = torch.nn.functional.interpolate(logits, size=image.size[::-1], mode="bilinear", align_corners=False)
-    seg = upsampled_logits.argmax(dim=1)[0]
-    # 배경(0)이 아닌 부분만 추출
-    mask = seg != 0
-    # 원본 이미지에 마스크 적용
-    masked_image = Image.fromarray(image).convert("RGBA")
-    data = masked_image.getdata()
-    newData = []
-    for i, item in enumerate(data):
-        if mask.flatten()[i]:
-            newData.append(item)
-        else:
-            newData.append((255, 255, 255, 0))
-    masked_image.putdata(newData)
-    return masked_image
+    output_image = remove(image)
+    return output_image
 
-# Gradio 인터페이스 정의
-iface = gr.Interface(fn=remove_background, inputs=gr.Image(type="pil"), outputs="image", title="Background Remover", description="Upload an image to remove the background.")
+iface = gr.Interface(
+    fn=remove_background,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=gr.outputs.Image(type="pil"),
+    title="배경 제거 도구",
+    description="이 도구는 이미지에서 배경을 제거합니다. 이미지를 업로드하면 배경이 제거된 이미지를 받을 수 있습니다."
+)
 
-# 앱 실행
 iface.launch()