Spaces:
Sleeping
Sleeping
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") | |
#url = "http://images.cocodataset.org/val2017/000000039769.jpg" | |
#image = Image.open(requests.get(url, stream=True).raw) | |
#st.success("Image open: Success") | |
#feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large-ade") | |
#model = DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade") | |
#st.success("Load model: Success") | |
#inputs = feature_extractor(images=image, return_tensors="pt") | |
#st.success("Feature extraction: Success") | |
#outputs = model(**inputs) | |
#logits = outputs.logits | |
#st.text(str(logits)) | |
#st.success("Success") |