File size: 2,808 Bytes
6d580b8
 
 
 
 
 
 
 
 
b2c1780
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6d580b8
 
 
 
 
 
 
 
 
 
 
9c9afb9
6d580b8
 
 
 
b2c1780
 
bcfb055
b2c1780
bcfb055
b2c1780
8cf7460
 
35d2c61
51cca02
32767fa
9c9afb9
35d2c61
 
 
 
 
 
 
 
 
bcfb055
 
8cf7460
2ab7dd8
bcfb055
8c84c61
6d580b8
 
 
 
 
 
 
8cf7460
 
6d580b8
 
 
39c6ef5
b2c1780
6d580b8
b2c1780
6d580b8
18ca939
b2c1780
6d580b8
b2c1780
 
6d580b8
bc19878
6d580b8
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
# Portrait Photo Generator App

# Imports
from PIL import Image, ImageFilter
import numpy as np
from transformers import pipeline
import gradio as gr
import os

models1 = [
    "facebook/detr-resnet-50-panoptic",
    "CIDAS/clipseg-rd64-refined",
    "facebook/maskformer-swin-large-ade",
    "nvidia/segformer-b1-finetuned-cityscapes-1024-1024",    
]
current_model = models1[0]

models=[
    pipeline("image-segmentation", model="facebook/detr-resnet-50-panoptic"),
    pipeline("image-segmentation", model="facebook/maskformer-swin-large-ade"),
    pipeline("image-segmentation", model="CIDAS/clipseg-rd64-refined"),
    pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024"),
]

#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, model_choice):
    image = img_resize(image)
    model = models[model_choice]
    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")
            model_name = gr.Dropdown(show_label=False, choices=[m for m in models], type="index", value=current_model, interactive=True)
        with gr.Column():
            gal1=gr.Gallery(type="filepath").style(grid=6)
    with gr.Row():
        with gr.Column():
                object_output = gr.Dropdown(label="Objects")

    
    image_input.change(get_seg, inputs=[image_input, model_name], outputs=[gal1,object_output])



app.launch()