Evaluate U-shape Transformer Model For Image Enhancement

[keenanjohnson]

Previously, we have been only evaluating the U Net architecture for training auto encoder models (https://en.wikipedia.org/wiki/U-Net)

However, there are several other promising papers like the U Shape Transformer Model here: https://github.com/LintaoPeng/U-shape_Transformer_for_Underwater_Image_Enhancement

We should evaluate this and try to compare them.

[keenanjohnson]

I have begun this work and done a few things:

1. I added support for the model and parameters into UIE program: git change
2. I trained a model using this network using 1024 patches. This network trained for 150 iterations. The time on a Mac studio with M3 processor and 96GB of shared memory was around 1 hour per iteration, so this took roughly 150 hours. I have posted this model to our hugging face account here: https://huggingface.co/Seattle-Aquarium/ushape_transformer_1024_150_iterations
3. I processed the usual test images with this model and upload the results to this folder: https://drive.google.com/drive/folders/1eiF4mKTlcDup2T8ikP7lDjWMblt4-1Zv?usp=drive_link

[keenanjohnson]

I don't have any super obvious improvements to my visual eye over the previous u net architecture, so I think we may have to wait until we have a standard way of quantifying the variance with coralnet toolbox before we can compare this rigorously to the u net models.

[zhrandell]

Thanks for this, @keenanjohnson! To my eye, I'd actually say that the U-shape is an improvement over the U-net. Specifically, it seems like (1) the colors are more faithful relative to the hand-edited "truth," and (2) the color distortions along deeper areas also seem to perform a bit better compared to the U-net.

Oh, and (3) I also don't see any lines or artifacts . . . would I be correct in guessing that your Mac M3 processor was able to handle processing (inference) on a full image simultaneously? Or does the inference with the U-shape still break up the input TIFFs into individual patches? If the latter, then kudos, as I don't see any odd lines or other artifacts (on the subset of images that I inspected closely).

As the results you linked above were obtained via training on the full Hugging Face dataset, I'm now especially intrigued to see how the U-shape architecture performs on the recently partitioned "homogenous lighting" image set, along the "variable depth / variable lighting" image set that Reid has started to experiment with (at #14). I'm uncertain whether or not those partitioned folders have been pushed to Hugging Face (rather, I think he's been running them locally), but if not, pushing those folder sets to Hugging Face is something we'll want to tackle this upcoming week.

[keenanjohnson]

This was processed on my laptop using the UIE GUI, which is a mac, but has a normal amount of RAM (16GB), so I'm confused why you all still see artifacts as I don't see them, which implies that either you aren't using the most update to date GUI or thee is a difference between windows and mac in the UIE gui.

I don't believe @reidtosa has published any new versions of the datasets with a partitioned section.

[keenanjohnson]

FYI @reidtosa I also wrote a basic guide for interacting with hugging face datasets here: https://github.com/Seattle-Aquarium/CCR_image_processing/wiki/Managing-Hugging-Face-Datasets

Let me know if you have other specific questions or things I can clarify there.

[reidtosa]

Very exciting update!

@keenanjohnson, I'm realizing now that we still have un-cropped test images stored in this repo. I'll make the swap so the outputs can better match our hand-edits.

That being said, I agree with Zach that the model produces a much more faithful match to our hand-edits in most cases. Though, I am noticing some blurry splotches across some of the images that I haven't seen or experienced before. Taking a look at 2024_10_08_10-09-29, for example, I find it strange that there is one high resolution spot just right of center surrounded by lots of blur. Usually the center ~50% of our images retain the most definition and detail, but here it almost looks like that portion of the image is out of focus relative to the small patch. Something to look out for in future tests. As for the artifacts, that seems to have resolved itself on your end but I am still experiencing it on the mapping station. I am running v0.3.0...

I haven't yet pushed the new datasets to Huggingface, so thank you for the guide, Keenan! I should have new sets up soon.

[m-h-williams]

I don't have any super obvious improvements to my visual eye over the previous u net architecture, so I think we may have to wait until we have a standard way of quantifying the variance with coralnet toolbox before we can compare this rigorously to the u net models.

I'm working on this as documented in issue #16

[m-h-williams]

@keenanjohnson I’ve added a new SOP and supporting Python scripts to this repo to streamline our workflow for evaluating UIE-processed imagery against hand-edited training images.

This update includes:

• A full SOP (linked here) outlining the step-by-step process for preparing datasets, updating annotation paths, matching dataset splits, training YOLO classification models, and comparing performance metrics.

• A new code/ folder containing three scripts detailed in the SOP:

  • update_annotation_paths.py — updates annotation JSON files to point to UIE image directories
  • match_dataset_split.py — ensures UIE and hand-edited datasets have identical train/val/test splits
  • ML_metrics_comparison.py — generates a clean comparison table of model performance metrics

These should make the workflow much smoother, more reproducible, and quicker for anyone on the team to run. Let me know if you have ideas, questions, or anything you'd like to tweak!

[reidtosa]

@keenanjohnson, overall, I've been very impressed by the performance our U-Shape Transformer models; however, I've also picked up on a consistent trend in the outputs. The same several images always develop "smudges"/blurry spots in the same location regardless of epoch or which trained U-Shape model we apply. Shown below are blurry spots in 2024_10_08_11-40-27 after applying the "all01" model at epochs 25 and 158 (trained on combined flat and complex image sets) and the "flat06" best model (trained on only the flat image set).

It does seem like these blurry areas worsen over time, but from what I can tell they start to occur early in development. Furthermore, many of them occur in the center portion of the image, which should retain the most detail.

Looking back, Keenan, we also noticed these spots in some of the initial testing images you shared for U-Shape.

2. I trained a model using this network using 1024 patches. This network trained for 150 iterations. The time on a Mac studio with M3 processor and 96GB of shared memory was around 1 hour per iteration, so this took roughly 150 hours. I have posted this model to our hugging face account here: https://huggingface.co/Seattle-Aquarium/ushape_transformer_1024_150_iterations
3. I processed the usual test images with this model and upload the results to this folder: https://drive.google.com/drive/folders/1eiF4mKTlcDup2T8ikP7lDjWMblt4-1Zv?usp=drive_link

Interestingly, a little over half of the images have no blur at all and look great. Any thoughts on whether these artifacts can be worked out somewhere down the line? Any clues as to why this happens only to a selection of the images? I'll upload the models above to Hugging Face in case you'd like to see for yourself.

[keenanjohnson]

Hmmmm that is quite interesting @reidtosa . I have to assume that this artifact is an output of the model itself, rather than the GUI or post processing where we stitch patches together since we don't see it in the other model type correct?

Why that happens is unclear to me. One theory is that there is something about those spots which is tricking the model into processing in some other way that is incorrect (perhaps it estimates depth incorrectly, etc), but that's really just conjecture I would say.

[keenanjohnson]

Hey all! I've had a few moments here and there to work on this further and I have some learnings to share and hoping to have more time to come to an actual conclusion. I shared some of this with Zach last week, but finally had a moment to document.

The main question we've all had is whether these blurry artifacts are some sort of bug in the blending algorithm that processes tiles when using the GUI or some inherent issue with the trained model. To help get to the bottom of this, I first edited a version of the GUI script to do no overlapping or blending of the tiles during processing. It just tiles up the image and processes each through the trained model. The results as an example are shown in this image where you can clearly see each tile (this also illustrates why the blending / overlap is required as the color differences tile to tile can very widely).

Now, if the blurriness issue was purely from the blending or overlap process the GUI does, we would expect to see all non blurry tiles in this zero overlap method. But as you can see, in the problem spots the whole tile is clearly blurry. The processing python script does no modification to the pixel colors coming out of the model, so I think this conclusively proves that the blurry patches are an artifact of the trained U-shape transformer model in some way.

Now what to do about that haha. I think there are a few possibilities:

1. The blurriness is because the model does not have sufficient context as to the surrounding area and that by increasing the size of image patches we train the model on, we could give the model enough information to fix these blurry spots.
2. This is an inherent flaw of the current u-shape transformer model and we need to go back to the more basic U-net or try implementing another model from one of the papers.

Option 1 above is easy enough to test, and I think the next step is to try and increase the patch size for training, which would either confirm or rule out number 1, leaving only option 2.

[zhrandell]

Thanks for this, @keenanjohnson! Super interesting to see the base patches prior to the blending phase - your conclusion about the blurry patches reflecting something trained within the U-shape model makes sense to me.

Both of your proposed next steps sound good - I'll be curious to see if patch size affects the blur.

The other possibility we've discussed is that perhaps the model is optimizing around some sort of metric / goal that (for whatever reason) nudges the U-shape towards blurry patches; I believe this is to some degree evidenced by the blurry patches getting worse as the epochs progress, such that in at least one model training session, I believe the final (best) epoch produced fully blurred out images. (@reidtosa would need to remind us about the particulars of that model training session).

In any case, I'm hoping us having access to some beefier compute in the new year will help us experiment with, e.g., larger patch sizes, different models (if warranted), etc.

Thanks again!