import gradio as gr
from PIL import Image

inimg = gr.Image()
outimg = gr.AnnotatedImage()
custom_html = '''
<center>
  <div style="overflow:hidden ; max-width: fit-content; margin-left: auto; margin-right: auto;">
      <div style="float: left; font-family: Arial; font-size: 25px; color: orange; margin: 10px;">Please</div>
        <div style="float: left">
            <a href="https://www.buymeacoffee.com/alloc7260">
                <img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=☕&slug=alloc7260&button_colour=FFDD00&font_colour=000000&font_family=Arial&outline_colour=000000&coffee_colour=ffffff" />
            </a>
        </div>
      <div style="float: left; font-family: Arial; font-size: 25px; color: orange; margin: 10px;">for upgrade to gpu instance</div>
  </div>
</center>
'''

def resize_image(image, max_width=1500):
    original_width, original_height = image.size
    aspect_ratio = original_width / original_height
    new_width = min(original_width, max_width)
    new_height = int(new_width / aspect_ratio)
    resized_image = image.resize((new_width, new_height), Image.LANCZOS)
    return resized_image

def mask(img):
  PIL_image = Image.fromarray(img.astype('uint8'), 'RGB')
  resimg = resize_image(PIL_image)
  import torch
  from transformers import pipeline
  with torch.inference_mode():
    generator =  pipeline("mask-generation", "facebook/sam-vit-base", points_per_batch = 64)
    outputs = generator(resimg, points_per_batch = 64)
    outputs_masks = outputs['masks']
    torch.cuda.empty_cache()
  return (resimg, [(outputs_masks[i], str(i + 1)) for i in range(len(outputs_masks))])

interface = gr.Blocks()

with interface:
  gr.Interface(mask, inimg, outimg)
  che = gr.HTML(custom_html)

interface.launch(show_api=False, debug=True)