File size: 1,099 Bytes
8b04a03
10cb35a
2aa7829
7899b33
 
 
8b04a03
2aa7829
10cb35a
 
7899b33
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10cb35a
 
 
 
2aa7829
 
 
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
import streamlit as st
from PIL import Image
from transformers import pipeline
import numpy as np
import cv2
import matplotlib.cm as cm

semantic_segmentation = pipeline("image-segmentation", "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png"])

def draw_masks_fromDict(image, results):
    masked_image = image.copy()
    
    colormap = cm.get_cmap('nipy_spectral')
    
    for i, result in enumerate(results):
        mask = np.array(result['mask'])  
        mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)  
        
        color = colormap(i / len(results))[:3] 
        color = tuple(int(c * 255) for c in color)  

        masked_image = np.where(mask, color, masked_image)

    masked_image = masked_image.astype(np.uint8)
    return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)


if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("")

    results = semantic_segmentation(image)
    st.json(results)