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") |