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 DPT')
st.write('The DPT model was proposed in Vision Transformers for Dense Prediction by René Ranftl, Alexey Bochkovskiy, Vladlen Koltun. DPT is a model that leverages the Vision Transformer (ViT) as backbone for dense prediction tasks like semantic segmentation and depth estimation.')
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)
	st.success("Load Image: Success")
	
	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")
		model = model.to(device)
		model.eval()
		
	st.success("Load model: Success")
	with st.spinner('Preparing image...'):
		# prepare the image for the model (aligned resize)
		feature_extractor_inference = DPTFeatureExtractor(do_random_crop=False, do_pad=False)
		pixel_values = feature_extractor_inference(image, return_tensors="pt").pixel_values.to(device)

	with st.spinner('Running inference...'):
		outputs = model(pixel_values=pixel_values)# logits are of shape (batch_size, num_labels, height/4, width/4)

	with st.spinner('Postprocessing...'):
		logits = outputs.logits.cpu()
		# First, rescale logits to original image size
		upsampled_logits = nn.functional.interpolate(logits,
		                size=image.shape[:-1], # (height, width)
		                mode='bilinear',
		                align_corners=False)
		# Second, apply argmax on the class dimension
		seg = upsampled_logits.argmax(dim=1)[0]
		color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3\
		all_labels = []
		for label, color in enumerate(palette):
		    color_seg[seg == label, :] = color
		    if label in seg:
		    	all_labels.append(id2label[label])
		# Convert to BGR
		color_seg = color_seg[..., ::-1]
	# Show image + mask
	img = np.array(image) * 0.5 + color_seg * 0.5
	img = img.astype(np.uint8)
	st.image(img, caption="Segmented Image")
	st.header("Predicted Labels")
	for idx, label in enumerate(all_labels):
		st.subheader(f'{idx+1}) {label}')	

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