added spinners for all tasks

app.py

@@ -20,33 +20,38 @@ if raw_image is not None:
 	image = Image.open(raw_image)
 	image = np.asarray(image)
 	
-	feature_extractor = SegformerFeatureExtractor(align=False, reduce_zero_label=False)
-	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-	model = SegformerForSemanticSegmentation.from_pretrained("deep-learning-analytics/segformer_semantic_segmentation", 
-															 ignore_mismatched_sizes=True,
-	                                                         num_labels=len(id2label), id2label=id2label, label2id=label2id,
-	                                                         reshape_last_stage=True)
-	model = model.to(device)
-	# prepare the image for the model (aligned resize)
-	feature_extractor_inference = SegformerFeatureExtractor(do_random_crop=False, do_pad=False)
-	pixel_values = feature_extractor_inference(image, return_tensors="pt").pixel_values.to(device)
-	model.eval()
+	with st.spinner('Loading Model...'):
+		feature_extractor = SegformerFeatureExtractor(align=False, reduce_zero_label=False)
+		device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+		model = SegformerForSemanticSegmentation.from_pretrained("deep-learning-analytics/segformer_semantic_segmentation", 
+																 ignore_mismatched_sizes=True,
+		                                                         num_labels=len(id2label), id2label=id2label, label2id=label2id,
+		                                                         reshape_last_stage=True)
+		model = model.to(device)
+		model.eval()
+	
+	with st.spinner('Preparing image...'):
+		# prepare the image for the model (aligned resize)
+		feature_extractor_inference = SegformerFeatureExtractor(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)
-	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\
-	for label, color in enumerate(palette):
-	    color_seg[seg == label, :] = color
-	# Convert to BGR
-	color_seg = color_seg[..., ::-1]
+
+	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\
+		for label, color in enumerate(palette):
+		    color_seg[seg == label, :] = color
+		# 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)