import gradio as gr
from transformers import pipeline
from PIL import Image

import librosa
import numpy as np
import torch

from transformers import AutoProcessor, AutoModelForCausalLM

ic_processor = AutoProcessor.from_pretrained("microsoft/git-base")
ic_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base")

def remove_background(image):
  pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
  pillow_mask = pipe(image, return_mask = True)
  pillow_image = pipe(image)

  return pillow_image, predict(image)


def predict(image):
    pixel_values = ic_processor(images=image, return_tensors="pt").pixel_values
    text_ids = ic_model.generate(pixel_values=pixel_values, max_length=50)
    text = ic_processor.batch_decode(text_ids, skip_special_tokens=True)[0]
    
    return text

theme = gr.themes.Glass(
    primary_hue="gray",
    secondary_hue="green",
    neutral_hue="gray",
)

with gr.Blocks(theme=theme) as app:

  gr.Markdown('# Publicação de Produto')
  gr.Markdown('Esta aplicação realiza de forma automática a preparação de um produto de e-commerce, removendo o fundo e gerando a descrição do produto a partir da imagem enviada.')

  with gr.Row():

    with gr.Column():
      gr.Markdown('## Selecione a imagem a ser tratada')
      imagem_in=gr.components.Image(type='pil')
    with gr.Column():
      gr.Markdown('## Imagem tratada')
      imagem_out=gr.components.Image(type='pil', format='png')
      gr.Markdown('## Descrição do anúncio da image, a partir da imagem enviada.')
      destino_texto=gr.components.Textbox(value='')

  botao=gr.Button('Gerar Anuncio')
  botao.click(
      fn=remove_background, 
      inputs=imagem_in, 
      outputs=[imagem_out, destino_texto]
  )

app.launch(share=True)