Update app.py

app.py

@@ -3,81 +3,58 @@ import torch
 import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
 import gradio as gr
-from gradio_imageslider import ImageSlider
 from briarmbg import BriaRMBG
 import PIL
 from PIL import Image
 from typing import Tuple
 
-
+# Carregar o modelo pré-treinado
 net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 net.to(device)
 
-    
+# Função para redimensionar a imagem para o tamanho que o modelo espera
 def redimensionar_imagem(imagem):
     imagem = imagem.convert('RGB')
     tamanho_entrada_modelo = (1024, 1024)
     imagem = imagem.resize(tamanho_entrada_modelo, Image.BILINEAR)
     return imagem
 
-
+# Função principal para processar a imagem
 def processar(imagem):
-
     # preparar entrada
     imagem_original = Image.fromarray(imagem)
-    w,h = tamanho_imagem_original = imagem_original.size
+    w, h = imagem_original.size
     imagem = redimensionar_imagem(imagem_original)
     im_np = np.array(imagem)
-    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
-    im_tensor = torch.unsqueeze(im_tensor,0)
-    im_tensor = torch.divide(im_tensor,255.0)
-    im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
+    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
+    im_tensor = torch.unsqueeze(im_tensor, 0)
+    im_tensor = torch.divide(im_tensor, 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()
+        im_tensor = im_tensor.cuda()
 
-    # inferência
+    # Inferência com o modelo
     resultado = net(im_tensor)
-    # pós-processamento
-    resultado = torch.squeeze(F.interpolate(resultado[0][0], size=(h,w), mode='bilinear') ,0)
+    
+    # Pós-processamento
+    resultado = torch.squeeze(F.interpolate(resultado[0][0], size=(h, w), mode='bilinear'), 0)
     ma = torch.max(resultado)
     mi = torch.min(resultado)
-    resultado = (resultado-mi)/(ma-mi)    
-    # imagem para pil
-    im_array = (resultado*255).cpu().data.numpy().astype(np.uint8)
+    resultado = (resultado - mi) / (ma - mi)
+    
+    # Convertendo o resultado para imagem PIL
+    im_array = (resultado * 255).cpu().data.numpy().astype(np.uint8)
     pil_im = Image.fromarray(np.squeeze(im_array))
-    # colar a máscara na imagem original
-    nova_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
+    
+    # Colando a máscara na imagem original
+    nova_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
     nova_im.paste(imagem_original, mask=pil_im)
 
     return nova_im
 
-
-# block = gr.Blocks().queue()
-
-# with block:
-#     gr.Markdown("## BRIA RMBG 1.4")
-#     gr.HTML('''
-#       <p style="margin-bottom: 10px; font-size: 94%">
-#         Esta é uma demonstração do BRIA RMBG 1.4 que utiliza
-#         <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">o modelo de matting de imagem BRIA RMBG-1.4</a> como base. 
-#       </p>
-#     ''')
-#     with gr.Row():
-#         with gr.Column():
-#             input_image = gr.Image(sources=None, type="pil") # None para upload, ctrl+v e webcam
-#             # input_image = gr.Image(sources=None, type="numpy") # None para upload, ctrl+v e webcam
-#             run_button = gr.Button(value="Executar")
-            
-#         with gr.Column():
-#             result_gallery = gr.Gallery(label='Resultado', show_label=False, elem_id="gallery", columns=[1], height='auto')
-#     ips = [input_image]
-#     run_button.click(fn=processar, inputs=ips, outputs=[result_gallery])
-
-# block.launch(debug = True)
-
-# block = gr.Blocks().queue()
-
+# Interface com Gradio
 gr.Markdown("## BRIA RMBG 1.4")
 gr.HTML('''
   <p style="margin-bottom: 10px; font-size: 94%">
@@ -85,10 +62,27 @@ gr.HTML('''
     <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">o modelo de matting de imagem BRIA RMBG-1.4</a> como base. 
   </p>
 ''')
+
+# Exemplos
 exemplos = [['./input.jpg'],]
-# output = ImageSlider(position=0.5,label='Imagem sem fundo', type="pil", show_download_button=True)
-# demo = gr.Interface(fn=processar,inputs="image", outputs=output, examples=exemplos, title=title, description=description)
-demo = gr.Interface(fn=processar,inputs="image", outputs="image", examples=exemplos)
 
-if __name__ == "__main__":
-    demo.launch(share=False)
+# Função para criar um botão de download
+def preparar_download(imagem_processada):
+    return gr.File.update(value=imagem_processada, label="Baixar Imagem Processada", visible=True)
+
+# Configurando a interface
+output = gr.Image(type="pil", label="Imagem Processada")
+download_output = gr.File(label="Baixar Imagem Processada", visible=False)  # Inicialmente invisível
+
+# Definindo a interface com inputs e outputs
+demo = gr.Interface(
+    fn=processar,
+    inputs=gr.Image(type="numpy", label="Carregar Imagem"),  # Traduzido "Image" para "Carregar Imagem"
+    outputs=[output, download_output],
+    examples=exemplos,
+    title="BRIA RMBG 1.4",
+    description="Remover o fundo da imagem usando o modelo BRIA RMBG 1.4.",
+)
+
+# Atualizando o botão de download após a imagem ser processada
+demo.launch(share=False)