|
import os |
|
import pandas as pd |
|
import numpy as np |
|
import torch |
|
from transformers import DPTFeatureExtractor, DPTForSemanticSegmentation |
|
from PIL import Image |
|
from torch import nn |
|
import requests |
|
import streamlit as st |
|
|
|
img_path = None |
|
st.title('Semantic Segmentation using Beit') |
|
file_upload = st.file_uploader('Raw Input Image') |
|
image_path = st.selectbox( |
|
'Choose any one image for inference', |
|
('Select image', 'image1.jpg', 'image2.jpg', 'image3.jpg')) |
|
|
|
if file_upload is None: |
|
raw_image = image_path |
|
else: |
|
raw_image = file_upload |
|
|
|
if raw_image != 'Select image': |
|
df = pd.read_csv('class_dict_seg.csv') |
|
classes = df['name'] |
|
palette = df[[' r', ' g', ' b']].values |
|
id2label = classes.to_dict() |
|
label2id = {v: k for k, v in id2label.items()} |
|
|
|
image = Image.open(raw_image) |
|
image = np.asarray(image) |
|
|
|
with st.spinner('Loading Model...'): |
|
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large-ade") |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
model = DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade",ignore_mismatched_sizes=True,num_labels=len(id2label), id2label=id2label, label2id=label2id,reshape_last_stage=True) |
|
model = model.to(device) |
|
model.eval() |
|
|
|
|
|
st.success("Success") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|