import gradio as gr
import pixellib
from pixellib.torchbackend.instance import instanceSegmentation
from PIL import Image
import cv2 

ins = instanceSegmentation()
ins.load_model("https://github.com/ayoolaolafenwa/PixelLib/releases/download/0.2.0/pointrend_resnet50.pkl")
target_classes = ins.select_target_classes(person=True)
#seg.segmentImage("sample2.jpg", segment_target_classes= target_classes, show_bboxes=True,  output_image_name="a.jpg") 


def seg(inp):
    out_box=[]
    results, output = ins.segmentImage(f"{inp}",
                                       #segment_target_classes= target_classes,
                                       show_bboxes=False,
                                       extract_segmented_objects= True,
                                       save_extracted_objects=True,
                                       output_image_name="output_image.jpg")
    #results, output = ins.segmentImage("image.jpg", show_bboxes=True, output_image_name="result.jpg")
    print (results)
    for image in results['extracted_objects']:
        out_box.append(Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)))
    return out_box
'''
gr.Interface(
    fn=seg,
    inputs=gr.Image(type='filepath'),
    outputs=gr.Gallery()
).launch()
'''
with gr.Blocks() as app:
    inp=gr.Image(type='filepath')
    btn=gr.Button()
    outp=gr.Gallery()
    btn.click(seg,inp,outp)
app.launch()