Since the mapping is “hard-coded” you could keep the model with the class outputs and later use the prediction to map it to the corresponding category.
However, another way would be to let the model learn to combine the classes to categories.
To do so you could add another layer, get the class predictions as its input, and output the categories.
This might allow the model to learn that e.g. 5 “medium predictions” of flat classes are a stronger signal than a single higher prediction of a non-flat class, assuming that the target is the flat category.