Rename main.py to app.py

app.py

@@ -0,0 +1,27 @@
+from fastapi import FastAPI, HTTPException
+import httpx
+from rembg import remove
+from io import BytesIO
+from fastapi.responses import StreamingResponse
+
+app = FastAPI()
+
+@app.get("/remove-background/")
+async def remove_background(image_url: str):
+    try:
+        # Baixa a imagem da URL fornecida
+        async with httpx.AsyncClient() as client:
+            response = await client.get(image_url)
+            response.raise_for_status()  # Verifica se a URL foi bem-sucedida
+            image_data = response.content
+        
+        # Remover o fundo da imagem usando o rembg
+        output_image = remove(image_data)
+
+        # Retorna a imagem com o fundo removido como resposta
+        return StreamingResponse(BytesIO(output_image), media_type="image/png")
+    
+    except httpx.RequestError as e:
+        raise HTTPException(status_code=400, detail="Erro ao baixar a imagem: " + str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="Erro ao processar a imagem: " + str(e))

main.py

@@ -1,89 +0,0 @@
-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 fastapi import FastAPI, HTTPException
-from PIL import Image
-import httpx
-import io
-import briarmbg  # Importando o modelo de remoção de fundo
-
-app = FastAPI()
-
-# Carregar modelo BRIA RMBG
-net = briarmbg.BriaRMBG()
-model_path = hf_hub_download("briaai/RMBG-1.4", "model.pth")
-
-if torch.cuda.is_available():
-    net.load_state_dict(torch.load(model_path))
-    net = net.cuda()
-else:
-    net.load_state_dict(torch.load(model_path, map_location="cpu"))
-
-net.eval()
-
-# Função para redimensionar a imagem antes de processar
-def resize_image(image):
-    image = image.convert("RGB")
-    model_input_size = (1024, 1024)
-    return image.resize(model_input_size, Image.BILINEAR)
-
-# Função para remover o fundo da imagem
-def remove_bg(image: Image.Image):
-    orig_image = image
-    w, h = orig_image.size
-    image = resize_image(orig_image)
-    
-    im_np = np.array(image)
-    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
-    im_tensor = torch.unsqueeze(im_tensor, 0) / 255.0
-    im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
-    
-    if torch.cuda.is_available():
-        im_tensor = im_tensor.cuda()
-
-    # Inference do modelo
-    result = net(im_tensor)[0][0]
-    result = torch.squeeze(F.interpolate(result, size=(h, w), mode="bilinear"), 0)
-
-    # Normalizar para 0-255
-    ma, mi = torch.max(result), torch.min(result)
-    result = (result - mi) / (ma - mi)
-    
-    im_array = (result * 255).cpu().data.numpy().astype(np.uint8)
-    pil_im = Image.fromarray(np.squeeze(im_array))
-
-    # Criar imagem sem fundo
-    new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
-    new_im.paste(orig_image, mask=pil_im)
-
-    return new_im
-
-# Função para baixar imagem de uma URL
-async def download_image(image_url: str) -> Image.Image:
-    async with httpx.AsyncClient() as client:
-        response = await client.get(image_url)
-        if response.status_code != 200:
-            raise HTTPException(status_code=400, detail="Erro ao baixar imagem")
-        
-        return Image.open(io.BytesIO(response.content))
-
-# Endpoint para remover fundo
-@app.get("/remove-bg/")
-async def remove_bg_from_url(image_url: str):
-    try:
-        image = await download_image(image_url)
-        output_image = remove_bg(image)
-
-        # Salvar a imagem temporariamente na memória
-        img_io = io.BytesIO()
-        output_image.save(img_io, format="PNG")
-        img_io.seek(0)
-
-        return {
-            "message": "Fundo removido com sucesso!",
-            "image": img_io.getvalue()
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))