Sm00thix
/

unet

@@ -13,7 +13,7 @@ pipeline_tag: image-segmentation
 # U-Net
 This repository contains an implementation of U-Net [[1]](#references). [unet.py](./unet.py) implements the class UNet. The implementation has been tested with PyTorch 2.7.1 and CUDA 12.6.
-![](./assets/unet_diagram.png)
 You can also load the U-Net from PyTorch Hub.
 ```python
@@ -56,7 +56,7 @@ If you use this U-Net implementation, please cite Engstrøm et al. [[5]](#refere
 ## Citation
 If you use the code shared in this repository, please cite this work: https://arxiv.org/abs/2504.14131. The U-Net implementation in this repository was used to generate pixel-wise fat predictions in an image of a pork belly.
-![](./assets/unet_flow.png)
 ## References

 # U-Net
 This repository contains an implementation of U-Net [[1]](#references). [unet.py](./unet.py) implements the class UNet. The implementation has been tested with PyTorch 2.7.1 and CUDA 12.6.
+![](./images/unet_diagram.png)
 You can also load the U-Net from PyTorch Hub.
 ```python
 ## Citation
 If you use the code shared in this repository, please cite this work: https://arxiv.org/abs/2504.14131. The U-Net implementation in this repository was used to generate pixel-wise fat predictions in an image of a pork belly.
+![](./images/unet_flow.png)
 ## References