How to load a SegFormer pytorch checkpointfile and visualize a prediction through it?

Question

I have trained a segformer model using the code in this article (courtesy of Roboflow)

https://blog.roboflow.com/how-to-train-segformer-on-a-custom-dataset-with-pytorch-lightning/'

I end up with a checkpoint file, and now I want a seperat code to load the best.ckpt (model), load an image (convert it to a tensor), then run the image through the model and visualize the predicted masks.

I've tried various tutorials, but to no avail. What would be the minimum code to achieve this?

As far as I've gathered, these should be the approximate steps

1. Load model
2. Load image
3. Convert image to tensor
4. Run the image as prediction through the model
5. Convert the resulting prediction to a new image
6. Overlay that image onto the original image to show masks on image

I've Run through several examples, but hit dead ends when loading my own .ckpt file.

I saw this: How to use a trained SegFormer model on a single image? but the resolution is inconclusive

I've put the code stated below together, but get the following errors when I try to load the model:

Missing key(s) in state_dict: "segformer.encoder.patch_embeddings.0.proj.weight", "segformer.encoder.patch_embeddings.0.proj.bias", "segformer.encoder.patch_embeddings.0.layer_norm.weight", "segformer.encoder.patch_embeddings.0.layer_norm.bias", "segformer.encoder.patch_embeddings.1.proj.weight", "segformer.encoder.patch_embeddings.1.proj.bias", "segformer.encoder.patch_embeddings.1.layer_norm.weight", "segformer.encoder.patch_embeddings.1.layer_norm.bias", "segformer.encoder.patch_embeddings.2.proj.weight", "segformer.encoder.patch_embeddings.2.proj.bias", "segformer.encoder.patch_embeddings.2.layer_norm.weight", "segformer.encoder.patch_embeddings.2.layer_norm.bias", "segformer.encoder.patch_embeddings.3.proj.weight", "segformer.encoder.patch_embeddings.3.proj.bias", "segformer.encoder.patch_embeddings.3.layer_norm.weight", "segformer.encoder.patch_embeddings.3.layer_norm.bias", "segformer.encoder.block.0.0.layer_norm_1.weight", "segformer.encoder.block.0.0.layer_norm_1.bias", "segformer.encoder.block.0.0.attention.self.query.weight", "segformer.encoder.block.0.0.attention.self.query.bias", "segformer.encoder.block.0.0.attention.self.key.weight", "segformer.encoder.block.0.0.attention.self.key.bias", "segformer.encoder.block.0.0.attention.self.value.weight", "segformer.encoder.block.0.0.attention.self.value.bias", "segformer.encoder.block.0.0.attention.self.sr.weight", "segformer.encoder.block.0.0.attention.self.sr.bias", "segformer.encoder.block.0.0.attention.self.layer_norm.weight", "segformer.encoder.block.0.0.attention.self.layer_norm.bias", "segformer.encoder.block.0.0.attention.output.dense.weight", "segformer.encoder.block.0.0.attention.output.dense.bias", "segformer.encoder.block.0.0.layer_norm_2.weight", "segformer.encoder.block.0.0.layer_norm_2.bias", "segformer.encoder.block.0.0.mlp.dense1.weight", "segformer.encoder.block.0.0.mlp.dense1.bias", "segformer.encoder.block.0.0.mlp.dwconv.dwconv.weight", "segformer.encoder.block.0.0.mlp.dwconv.dwconv.bias", "segformer.encoder.block.0.0.mlp.dense2.weight", "segformer.encoder.block.0.0.mlp.dense2.bias", "segformer.encoder.block.0.1.layer_norm_1.weight", "segformer.encoder.block.0.1.layer_norm_1.bias", "segformer.encoder.block.0.1.attention.self.query.weight", "segformer.encoder.block.0.1.attention.self.query.bias", "segformer.encoder.block.0.1.attention.self.key.weight", "segformer.encoder.block.0.1.attention.self.key.bias", "segformer.encoder.block.0.1.attention.self.value.weight", "segformer.encoder.block.0.1.attention.self.value.bias", "segformer.encoder.block.0.1.attention.self.sr.weight", "segformer.encoder.block.0.1.attention.self.sr.bias", "segformer.encoder.block.0.1.attention.self.layer_norm.weight", "segformer.encoder.block.0.1.attention.self.layer_norm.bias", "segformer.encoder.block.0.1.attention.output.dense.weight", "segformer.encoder.block.0.1.attention.output.dense.bias", "segformer.encoder.block.0.1.layer_norm_2.weight", "segformer.encoder.block.0.1.layer_norm_2.bias", "segformer.encoder.block.0.1.mlp.dense1.weight", "segformer.encoder.block.0.1.mlp.dense1.bias", "segformer.encoder.block.0.1.mlp.dwconv.dwconv.weight", "segformer.encoder.block.0.1.mlp.dwconv.dwconv.bias", "segformer.encoder.block.0.1.mlp.dense2.weight", "segformer.encoder.block.0.1.mlp.dense2.bias", "segformer.encoder.block.1.0.layer_norm_1.weight", "segformer.encoder.block.1.0.layer_norm_1.bias", "segformer.encoder.block.1.0.attention.self.query.weight", "segformer.encoder.block.1.0.attention.self.query.bias", "segformer.encoder.block.1.0.attention.self.key.weight", "segformer.encoder.block.1.0.attention.self.key.bias", "segformer.encoder.block.1.0.attention.self.value.weight", "segformer.encoder.block.1.0.attention.self.value.bias", "segformer.encoder.block.1.0.attention.self.sr.weight", "segformer.encoder.block.1.0.attention.self.sr.bias", "segformer.encoder.block.1.0.attention.self.layer_norm.weight", "segformer.encoder.block.1.0.attention.self.layer_norm.bias", "segformer.encoder.block.1.0.attention.output.dense.weight", "segformer.encoder.block.1.0.attention.output.dense.bias", "segformer.encoder.block.1.0.layer_norm_2.weight", "segformer.encoder.block.1.0.layer_norm_2.bias", "segformer.encoder.block.1.0.mlp.dense1.weight", "segformer.encoder.block.1.0.mlp.dense1.bias", "segformer.encoder.block.1.0.mlp.dwconv.dwconv.weight", "segformer.encoder.block.1.0.mlp.dwconv.dwconv.bias", "segformer.encoder.block.1.0.mlp.dense2.weight", "segformer.encoder.block.1.0.mlp.dense2.bias", "segformer.encoder.block.1.1.layer_norm_1.weight", "segformer.encoder.block.1.1.layer_norm_1.bias", "segformer.encoder.block.1.1.attention.self.query.weight", "segformer.encoder.block.1.1.attention.self.query.bias", "segformer.encoder.block.1.1.attention.self.key.weight", "segformer.encoder.block.1.1.attention.self.key.bias", "segformer.encoder.block.1.1.attention.self.value.weight", "segformer.encoder.block.1.1.attention.self.value.bias", "segformer.encoder.block.1.1.attention.self.sr.weight", "segformer.encoder.block.1.1.attention.self.sr.bias", "segformer.encoder.block.1.1.attention.self.layer_norm.weight", "segformer.encoder.block.1.1.attention.self.layer_norm.bias", "segformer.encoder.block.1.1.attention.output.dense.weight", "segformer.encoder.block.1.1.attention.output.dense.bias", "segformer.encoder.block.1.1.layer_norm_2.weight", "segformer.encoder.block.1.1.layer_norm_2.bias", "segformer.encoder.block.1.1.mlp.dense1.weight", "segformer.encoder.block.1.1.mlp.dense1.bias", "segformer.encoder.block.1.1.mlp.dwconv.dwconv.weight", "segformer.encoder.block.1.1.mlp.dwconv.dwconv.bias", "segformer.encoder.block.1.1.mlp.dense2.weight", "segformer.encoder.block.1.1.mlp.dense2.bias", "segformer.encoder.block.2.0.layer_norm_1.weight", "segformer.encoder.block.2.0.layer_norm_1.bias", "segformer.encoder.block.2.0.attention.self.query.weight", "segformer.encoder.block.2.0.attention.self.query.bias", "segformer.encoder.block.2.0.attention.self.key.weight", "segformer.encoder.block.2.0.attention.self.key.bias", "segformer.encoder.block.2.0.attention.self.value.weight", "segformer.encoder.block.2.0.attention.self.value.bias", "segformer.encoder.block.2.0.attention.self.sr.weight", "segformer.encoder.block.2.0.attention.self.sr.bias", "segformer.encoder.block.2.0.attention.self.layer_norm.weight", "segformer.encoder.block.2.0.attention.self.layer_norm.bias", "segformer.encoder.block.2.0.attention.output.dense.weight", "segformer.encoder.block.2.0.attention.output.dense.bias", "segformer.encoder.block.2.0.layer_norm_2.weight", "segformer.encoder.block.2.0.layer_norm_2.bias", "segformer.encoder.block.2.0.mlp.dense1.weight", "segformer.encoder.block.2.0.mlp.dense1.bias", "segformer.encoder.block.2.0.mlp.dwconv.dwconv.weight", "segformer.encoder.block.2.0.mlp.dwconv.dwconv.bias", "segformer.encoder.block.2.0.mlp.dense2.weight", "segformer.encoder.block.2.0.mlp.dense2.bias", "segformer.encoder.block.2.1.layer_norm_1.weight", "segformer.encoder.block.2.1.layer_norm_1.bias", "segformer.encoder.block.2.1.attention.self.query.weight", "segformer.encoder.block.2.1.attention.self.query.bias", "segformer.encoder.block.2.1.attention.self.key.weight", "segformer.encoder.block.2.1.attention.self.key.bias", "segformer.encoder.block.2.1.attention.self.value.weight", "segformer.encoder.block.2.1.attention.self.value.bias", "segformer.encoder.block.2.1.attention.self.sr.weight", "segformer.encoder.block.2.1.attention.self.sr.bias", "segformer.encoder.block.2.1.attention.self.layer_norm.weight", "segformer.encoder.block.2.1.attention.self.layer_norm.bias", "segformer.encoder.block.2.1.attention.output.dense.weight", "segformer.encoder.block.2.1.attention.output.dense.bias", "segformer.encoder.block.2.1.layer_norm_2.weight", "segformer.encoder.block.2.1.layer_norm_2.bias", "segformer.encoder.block.2.1.mlp.dense1.weight", "segformer.encoder.block.2.1.mlp.dense1.bias", "segformer.encoder.block.2.1.mlp.dwconv.dwconv.weight", "segformer.encoder.block.2.1.mlp.dwconv.dwconv.bias", "segformer.encoder.block.2.1.mlp.dense2.weight", "segformer.encoder.block.2.1.mlp.dense2.bias", "segformer.encoder.block.3.0.layer_norm_1.weight", "segformer.encoder.block.3.0.layer_norm_1.bias", "segformer.encoder.block.3.0.attention.self.query.weight", "segformer.encoder.block.3.0.attention.self.query.bias", "segformer.encoder.block.3.0.attention.self.key.weight", "segformer.encoder.block.3.0.attention.self.key.bias", "segformer.encoder.block.3.0.attention.self.value.weight", "segformer.encoder.block.3.0.attention.self.value.bias", "segformer.encoder.block.3.0.attention.output.dense.weight", "segformer.encoder.block.3.0.attention.output.dense.bias", "segformer.encoder.block.3.0.layer_norm_2.weight", "segformer.encoder.block.3.0.layer_norm_2.bias", "segformer.encoder.block.3.0.mlp.dense1.weight", "segformer.encoder.block.3.0.mlp.dense1.bias", "segformer.encoder.block.3.0.mlp.dwconv.dwconv.weight", "segformer.encoder.block.3.0.mlp.dwconv.dwconv.bias", "segformer.encoder.block.3.0.mlp.dense2.weight", "segformer.encoder.block.3.0.mlp.dense2.bias", "segformer.encoder.block.3.1.layer_norm_1.weight", "segformer.encoder.block.3.1.layer_norm_1.bias", "segformer.encoder.block.3.1.attention.self.query.weight", "segformer.encoder.block.3.1.attention.self.query.bias", "segformer.encoder.block.3.1.attention.self.key.weight", "segformer.encoder.block.3.1.attention.self.key.bias", "segformer.encoder.block.3.1.attention.self.value.weight", "segformer.encoder.block.3.1.attention.self.value.bias", "segformer.encoder.block.3.1.attention.output.dense.weight", "segformer.encoder.block.3.1.attention.output.dense.bias", "segformer.encoder.block.3.1.layer_norm_2.weight", "segformer.encoder.block.3.1.layer_norm_2.bias", "segformer.encoder.block.3.1.mlp.dense1.weight", "segformer.encoder.block.3.1.mlp.dense1.bias", "segformer.encoder.block.3.1.mlp.dwconv.dwconv.weight", "segformer.encoder.block.3.1.mlp.dwconv.dwconv.bias", "segformer.encoder.block.3.1.mlp.dense2.weight", "segformer.encoder.block.3.1.mlp.dense2.bias", "segformer.encoder.layer_norm.0.weight", "segformer.encoder.layer_norm.0.bias", "segformer.encoder.layer_norm.1.weight", "segformer.encoder.layer_norm.1.bias", "segformer.encoder.layer_norm.2.weight", "segformer.encoder.layer_norm.2.bias", "segformer.encoder.layer_norm.3.weight", "segformer.encoder.layer_norm.3.bias", "decode_head.linear_c.0.proj.weight", "decode_head.linear_c.0.proj.bias", "decode_head.linear_c.1.proj.weight", "decode_head.linear_c.1.proj.bias", "decode_head.linear_c.2.proj.weight", "decode_head.linear_c.2.proj.bias", "decode_head.linear_c.3.proj.weight", "decode_head.linear_c.3.proj.bias", "decode_head.linear_fuse.weight", "decode_head.batch_norm.weight", "decode_head.batch_norm.bias", "decode_head.batch_norm.running_mean", "decode_head.batch_norm.running_var", "decode_head.classifier.weight", "decode_head.classifier.bias". 
    Unexpected key(s) in state_dict: "model.segformer.encoder.patch_embeddings.0.proj.weight", "model.segformer.encoder.patch_embeddings.0.proj.bias", "model.segformer.encoder.patch_embeddings.0.layer_norm.weight", "model.segformer.encoder.patch_embeddings.0.layer_norm.bias", "model.segformer.encoder.patch_embeddings.1.proj.weight", "model.segformer.encoder.patch_embeddings.1.proj.bias", "model.segformer.encoder.patch_embeddings.1.layer_norm.weight", "model.segformer.encoder.patch_embeddings.1.layer_norm.bias", "model.segformer.encoder.patch_embeddings.2.proj.weight", "model.segformer.encoder.patch_embeddings.2.proj.bias", "model.segformer.encoder.patch_embeddings.2.layer_norm.weight", "model.segformer.encoder.patch_embeddings.2.layer_norm.bias", "model.segformer.encoder.patch_embeddings.3.proj.weight", "model.segformer.encoder.patch_embeddings.3.proj.bias", "model.segformer.encoder.patch_embeddings.3.layer_norm.weight", "model.segformer.encoder.patch_embeddings.3.layer_norm.bias", "model.segformer.encoder.block.0.0.layer_norm_1.weight", "model.segformer.encoder.block.0.0.layer_norm_1.bias", "model.segformer.encoder.block.0.0.attention.self.query.weight", "model.segformer.encoder.block.0.0.attention.self.query.bias", "model.segformer.encoder.block.0.0.attention.self.key.weight", "model.segformer.encoder.block.0.0.attention.self.key.bias", "model.segformer.encoder.block.0.0.attention.self.value.weight", "model.segformer.encoder.block.0.0.attention.self.value.bias", "model.segformer.encoder.block.0.0.attention.self.sr.weight", "model.segformer.encoder.block.0.0.attention.self.sr.bias", "model.segformer.encoder.block.0.0.attention.self.layer_norm.weight", "model.segformer.encoder.block.0.0.attention.self.layer_norm.bias", "model.segformer.encoder.block.0.0.attention.output.dense.weight", "model.segformer.encoder.block.0.0.attention.output.dense.bias", "model.segformer.encoder.block.0.0.layer_norm_2.weight", "model.segformer.encoder.block.0.0.layer_norm_2.bias", "model.segformer.encoder.block.0.0.mlp.dense1.weight", "model.segformer.encoder.block.0.0.mlp.dense1.bias", "model.segformer.encoder.block.0.0.mlp.dwconv.dwconv.weight", "model.segformer.encoder.block.0.0.mlp.dwconv.dwconv.bias", "model.segformer.encoder.block.0.0.mlp.dense2.weight", "model.segformer.encoder.block.0.0.mlp.dense2.bias", "model.segformer.encoder.block.0.1.layer_norm_1.weight", "model.segformer.encoder.block.0.1.layer_norm_1.bias", "model.segformer.encoder.block.0.1.attention.self.query.weight", "model.segformer.encoder.block.0.1.attention.self.query.bias", "model.segformer.encoder.block.0.1.attention.self.key.weight", "model.segformer.encoder.block.0.1.attention.self.key.bias", "model.segformer.encoder.block.0.1.attention.self.value.weight", "model.segformer.encoder.block.0.1.attention.self.value.bias", "model.segformer.encoder.block.0.1.attention.self.sr.weight", "model.segformer.encoder.block.0.1.attention.self.sr.bias", "model.segformer.encoder.block.0.1.attention.self.layer_norm.weight", "model.segformer.encoder.block.0.1.attention.self.layer_norm.bias", "model.segformer.encoder.block.0.1.attention.output.dense.weight", "model.segformer.encoder.block.0.1.attention.output.dense.bias", "model.segformer.encoder.block.0.1.layer_norm_2.weight", "model.segformer.encoder.block.0.1.layer_norm_2.bias", "model.segformer.encoder.block.0.1.mlp.dense1.weight", "model.segformer.encoder.block.0.1.mlp.dense1.bias", "model.segformer.encoder.block.0.1.mlp.dwconv.dwconv.weight", "model.segformer.encoder.block.0.1.mlp.dwconv.dwconv.bias", "model.segformer.encoder.block.0.1.mlp.dense2.weight", "model.segformer.encoder.block.0.1.mlp.dense2.bias", "model.segformer.encoder.block.1.0.layer_norm_1.weight", "model.segformer.encoder.block.1.0.layer_norm_1.bias", "model.segformer.encoder.block.1.0.attention.self.query.weight", "model.segformer.encoder.block.1.0.attention.self.query.bias", "model.segformer.encoder.block.1.0.attention.self.key.weight", "model.segformer.encoder.block.1.0.attention.self.key.bias", "model.segformer.encoder.block.1.0.attention.self.value.weight", "model.segformer.encoder.block.1.0.attention.self.value.bias", "model.segformer.encoder.block.1.0.attention.self.sr.weight", "model.segformer.encoder.block.1.0.attention.self.sr.bias", "model.segformer.encoder.block.1.0.attention.self.layer_norm.weight", "model.segformer.encoder.block.1.0.attention.self.layer_norm.bias", "model.segformer.encoder.block.1.0.attention.output.dense.weight", "model.segformer.encoder.block.1.0.attention.output.dense.bias", "model.segformer.encoder.block.1.0.layer_norm_2.weight", "model.segformer.encoder.block.1.0.layer_norm_2.bias", "model.segformer.encoder.block.1.0.mlp.dense1.weight", "model.segformer.encoder.block.1.0.mlp.dense1.bias", "model.segformer.encoder.block.1.0.mlp.dwconv.dwconv.weight", "model.segformer.encoder.block.1.0.mlp.dwconv.dwconv.bias", "model.segformer.encoder.block.1.0.mlp.dense2.weight", "model.segformer.encoder.block.1.0.mlp.dense2.bias", "model.segformer.encoder.block.1.1.layer_norm_1.weight", "model.segformer.encoder.block.1.1.layer_norm_1.bias", "model.segformer.encoder.block.1.1.attention.self.query.weight", "model.segformer.encoder.block.1.1.attention.self.query.bias", "model.segformer.encoder.block.1.1.attention.self.key.weight", "model.segformer.encoder.block.1.1.attention.self.key.bias", "model.segformer.encoder.block.1.1.attention.self.value.weight", "model.segformer.encoder.block.1.1.attention.self.value.bias", "model.segformer.encoder.block.1.1.attention.self.sr.weight", "model.segformer.encoder.block.1.1.attention.self.sr.bias", "model.segformer.encoder.block.1.1.attention.self.layer_norm.weight", "model.segformer.encoder.block.1.1.attention.self.layer_norm.bias", "model.segformer.encoder.block.1.1.attention.output.dense.weight", "model.segformer.encoder.block.1.1.attention.output.dense.bias", "model.segformer.encoder.block.1.1.layer_norm_2.weight", "model.segformer.encoder.block.1.1.layer_norm_2.bias", "model.segformer.encoder.block.1.1.mlp.dense1.weight", "model.segformer.encoder.block.1.1.mlp.dense1.bias", "model.segformer.encoder.block.1.1.mlp.dwconv.dwconv.weight", "model.segformer.encoder.block.1.1.mlp.dwconv.dwconv.bias", "model.segformer.encoder.block.1.1.mlp.dense2.weight", "model.segformer.encoder.block.1.1.mlp.dense2.bias", "model.segformer.encoder.block.2.0.layer_norm_1.weight", "model.segformer.encoder.block.2.0.layer_norm_1.bias", "model.segformer.encoder.block.2.0.attention.self.query.weight", "model.segformer.encoder.block.2.0.attention.self.query.bias", "model.segformer.encoder.block.2.0.attention.self.key.weight", "model.segformer.encoder.block.2.0.attention.self.key.bias", "model.segformer.encoder.block.2.0.attention.self.value.weight", "model.segformer.encoder.block.2.0.attention.self.value.bias", "model.segformer.encoder.block.2.0.attention.self.sr.weight", "model.segformer.encoder.block.2.0.attention.self.sr.bias", "model.segformer.encoder.block.2.0.attention.self.layer_norm.weight", "model.segformer.encoder.block.2.0.attention.self.layer_norm.bias", "model.segformer.encoder.block.2.0.attention.output.dense.weight", "model.segformer.encoder.block.2.0.attention.output.dense.bias", "model.segformer.encoder.block.2.0.layer_norm_2.weight", "model.segformer.encoder.block.2.0.layer_norm_2.bias", "model.segformer.encoder.block.2.0.mlp.dense1.weight", "model.segformer.encoder.block.2.0.mlp.dense1.bias", "model.segformer.encoder.block.2.0.mlp.dwconv.dwconv.weight", "model.segformer.encoder.block.2.0.mlp.dwconv.dwconv.bias", "model.segformer.encoder.block.2.0.mlp.dense2.weight", "model.segformer.encoder.block.2.0.mlp.dense2.bias", "model.segformer.encoder.block.2.1.layer_norm_1.weight", "model.segformer.encoder.block.2.1.layer_norm_1.bias", "model.segformer.encoder.block.2.1.attention.self.query.weight", "model.segformer.encoder.block.2.1.attention.self.query.bias", "model.segformer.encoder.block.2.1.attention.self.key.weight", "model.segformer.encoder.block.2.1.attention.self.key.bias", "model.segformer.encoder.block.2.1.attention.self.value.weight", "model.segformer.encoder.block.2.1.attention.self.value.bias", "model.segformer.encoder.block.2.1.attention.self.sr.weight", "model.segformer.encoder.block.2.1.attention.self.sr.bias", "model.segformer.encoder.block.2.1.attention.self.layer_norm.weight", "model.segformer.encoder.block.2.1.attention.self.layer_norm.bias", "model.segformer.encoder.block.2.1.attention.output.dense.weight", "model.segformer.encoder.block.2.1.attention.output.dense.bias", "model.segformer.encoder.block.2.1.layer_norm_2.weight", "model.segformer.encoder.block.2.1.layer_norm_2.bias", "model.segformer.encoder.block.2.1.mlp.dense1.weight", "model.segformer.encoder.block.2.1.mlp.dense1.bias", "model.segformer.encoder.block.2.1.mlp.dwconv.dwconv.weight", "model.segformer.encoder.block.2.1.mlp.dwconv.dwconv.bias", "model.segformer.encoder.block.2.1.mlp.dense2.weight", "model.segformer.encoder.block.2.1.mlp.dense2.bias", "model.segformer.encoder.block.3.0.layer_norm_1.weight", "model.segformer.encoder.block.3.0.layer_norm_1.bias", "model.segformer.encoder.block.3.0.attention.self.query.weight", "model.segformer.encoder.block.3.0.attention.self.query.bias", "model.segformer.encoder.block.3.0.attention.self.key.weight", "model.segformer.encoder.block.3.0.attention.self.key.bias", "model.segformer.encoder.block.3.0.attention.self.value.weight", "model.segformer.encoder.block.3.0.attention.self.value.bias", "model.segformer.encoder.block.3.0.attention.output.dense.weight", "model.segformer.encoder.block.3.0.attention.output.dense.bias", "model.segformer.encoder.block.3.0.layer_norm_2.weight", "model.segformer.encoder.block.3.0.layer_norm_2.bias", "model.segformer.encoder.block.3.0.mlp.dense1.weight", "model.segformer.encoder.block.3.0.mlp.dense1.bias", "model.segformer.encoder.block.3.0.mlp.dwconv.dwconv.weight", "model.segformer.encoder.block.3.0.mlp.dwconv.dwconv.bias", "model.segformer.encoder.block.3.0.mlp.dense2.weight", "model.segformer.encoder.block.3.0.mlp.dense2.bias", "model.segformer.encoder.block.3.1.layer_norm_1.weight", "model.segformer.encoder.block.3.1.layer_norm_1.bias", "model.segformer.encoder.block.3.1.attention.self.query.weight", "model.segformer.encoder.block.3.1.attention.self.query.bias", "model.segformer.encoder.block.3.1.attention.self.key.weight", "model.segformer.encoder.block.3.1.attention.self.key.bias", "model.segformer.encoder.block.3.1.attention.self.value.weight", "model.segformer.encoder.block.3.1.attention.self.value.bias", "model.segformer.encoder.block.3.1.attention.output.dense.weight", "model.segformer.encoder.block.3.1.attention.output.dense.bias", "model.segformer.encoder.block.3.1.layer_norm_2.weight", "model.segformer.encoder.block.3.1.layer_norm_2.bias", "model.segformer.encoder.block.3.1.mlp.dense1.weight", "model.segformer.encoder.block.3.1.mlp.dense1.bias", "model.segformer.encoder.block.3.1.mlp.dwconv.dwconv.weight", "model.segformer.encoder.block.3.1.mlp.dwconv.dwconv.bias", "model.segformer.encoder.block.3.1.mlp.dense2.weight", "model.segformer.encoder.block.3.1.mlp.dense2.bias", "model.segformer.encoder.layer_norm.0.weight", "model.segformer.encoder.layer_norm.0.bias", "model.segformer.encoder.layer_norm.1.weight", "model.segformer.encoder.layer_norm.1.bias", "model.segformer.encoder.layer_norm.2.weight", "model.segformer.encoder.layer_norm.2.bias", "model.segformer.encoder.layer_norm.3.weight", "model.segformer.encoder.layer_norm.3.bias", "model.decode_head.linear_c.0.proj.weight", "model.decode_head.linear_c.0.proj.bias", "model.decode_head.linear_c.1.proj.weight", "model.decode_head.linear_c.1.proj.bias", "model.decode_head.linear_c.2.proj.weight", "model.decode_head.linear_c.2.proj.bias", "model.decode_head.linear_c.3.proj.weight", "model.decode_head.linear_c.3.proj.bias", "model.decode_head.linear_fuse.weight", "model.decode_head.batch_norm.weight", "model.decode_head.batch_norm.bias", "model.decode_head.batch_norm.running_mean", "model.decode_head.batch_norm.running_var", "model.decode_head.batch_norm.num_batches_tracked", "model.decode_head.classifier.weight", "model.decode_head.classifier.bias". 






import os
import pandas as pd
import numpy as np
import torch
from PIL import Image
from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor, SegformerConfig
from torch import nn
import torch




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("1.png").convert("RGB")
image = np.asarray(image)


feature_extractor = SegformerFeatureExtractor(align=False, reduce_zero_label=False)

checkpoint = torch.load("best.ckpt",map_location="cpu")
state_dict = checkpoint["state_dict"]

configuration = SegformerConfig()
model = SegformerForSemanticSegmentation(configuration)
model.load_state_dict(state_dict)
model.eval()



device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = model.to(device)
model.eval()

feature_extractor_inference = SegformerFeatureExtractor(do_random_crop=False, do_pad=False)
pixel_values = feature_extractor_inference(image, return_tensors="pt").pixel_values.to(device)


outputs = model(pixel_values)  # logits are of shape (batch_size, num_labels, height/4, width/4)

logits = outputs.logits
# 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)