Update app.py

app.py

@@ -1,20 +1,57 @@
+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
 
-url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-image = Image.open(requests.get(url, stream=True).raw)
-st.success("Image open: Success")
+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'))
 
-feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large-ade")
-model = DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade")
-st.success("Load model: Success")
+if file_upload is None:
+	raw_image = image_path
+else:
+	raw_image = file_upload
 
-inputs = feature_extractor(images=image, return_tensors="pt")
-st.success("Feature extraction: Success")
+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()}
 
-outputs = model(**inputs)
-logits = outputs.logits
+	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")
+#st.success("Success")