Image_Dissector / app.py
Omnibus's picture
Update app.py
39c6ef5
raw
history blame
2.08 kB
# Portrait Photo Generator App
# Imports
from PIL import Image, ImageFilter
import numpy as np
from transformers import pipeline
import gradio as gr
import os
model = pipeline("image-segmentation", model="facebook/detr-resnet-50-panoptic")
pred = []
def img_resize(image):
width = 1280
width_percent = (width / float(image.size[0]))
height = int((float(image.size[1]) * float(width_percent)))
return image.resize((width, height))
def image_objects(image):
global pred
image = img_resize(image)
pred = model(image)
pred_object_list = [str(i)+'_'+x['label'] for i, x in enumerate(pred)]
return gr.Dropdown.update(choices = pred_object_list, interactive = True)
def get_seg(image):
image = img_resize(image)
pred = model(image)
pred_object_list = [str(i)+'_'+x['label'] for i, x in enumerate(pred)]
seg_box=[]
for i in range(len(pred)):
#object_number = int(object.split('_')[0])
mask_array = np.asarray(pred[i]['mask'])/255
image_array = np.asarray(image)
mask_array_three_channel = np.zeros_like(image_array)
mask_array_three_channel[:,:,0] = mask_array
mask_array_three_channel[:,:,1] = mask_array
mask_array_three_channel[:,:,2] = mask_array
segmented_image = image_array*mask_array_three_channel
seg_out=segmented_image.astype(np.uint8)
seg_box.append(seg_out)
return(seg_box,gr.Dropdown.update(choices = pred_object_list, interactive = True))
app = gr.Blocks()
with app:
gr.Markdown(
"""
## Image Dissector
""")
with gr.Row():
with gr.Column():
image_input = gr.Image(label="Input Image",type="pil")
with gr.Row():
with gr.Column():
object_output = gr.Dropdown(label="Objects")
with gr.Row():
with gr.Column():
gal1=gr.Gallery(type="filepath").style(grid=4)
image_input.change(get_seg, inputs=[image_input], outputs=[gal1,object_output])
app.launch()