displaying predicted classes

app.py

@@ -48,11 +48,17 @@ if raw_image is not None:
 		# 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\
+		predicted_labels = []
 		for label, color in enumerate(palette):
 		    color_seg[seg == label, :] = color
+		    predicted_labels.append(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.image(img, caption="Segmented Image")
+
+	st.header('Predicted Labels')
+	for idx, label in enumerate(predicted_labels):
+		st.subheader(f'{idx+1}) {label}')