from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import gradio as gr
from PIL import Image
import cv2
import numpy as np
import torch

path = hf_hub_download("Bingsu/adetailer", "person_yolov8s-seg.pt")
model = YOLO(path)



def infer(img: Image.Image):
    img = "https://farm5.staticflickr.com/4139/4887614566_6b57ec4422_z.jpg"
    output = model(img)
    pred = output[0].plot()
    pred = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
    pred = Image.fromarray(pred)
    return pred


inputs = gr.inputs.Image(type='pil', label="Original Image")
outputs = [
           # gr.outputs.Image(type="numpy",label="Mask").style(height=600),
           gr.outputs.Image(type="numpy",label="Mask"),
        #    gr.outputs.Image(type="numpy",label="Foreground"),
           ]

examples = [
  ['fox.jpg'],
  ['parrot.jpg']
]

gr.Interface(infer, inputs, outputs, title=f"anime seg", allow_flagging=False).launch(server_name=f"10.17.1.16", server_port=6019)

if __name__ == "__main__":
    a = 1