Hi,
I’m quite new to ML and I’m trying to modify my current U-net image segmentation model (takes RGB input returns binary classification map) to give RGB output. I have made changes to the number of output channels, changed the loss function to cross entropy loss and removed thresholding as I want a continuous brightness level for the color channels instead of having just 0 or 1. The output images seem gray. All pixel values are the same in a channel. Now, in my R and B channel all values are 220 and in G channel all values are 0. What should I still change?
Hi Veera!
Let me speculate that instead of having your U-Net model perform
semantic segmentation (e.g., foreground vs. background), you would
like your model to do something like denoise or deblur your input image.
In such a case, you don’t want the pixels of your output image to be
predicted class labels, but rather predicted RGB pixel values.
This is incorrect for what I assume is your use case.
To train a model to predict continuous values (that don’t somehow
represent probabilities) you would typically use MSELoss.
So if your training data consists of input images with noise added to
them and target images that are the original images without the added
noise, you would want the (per-pixel) output of your U-Net to have
three channels (the RGB) and you would compute the MSELoss of
the output image with respect to the (three-channel, RGB) target image.
This is simply saying that you want the denoised image produced by
the model to be as close as possible to the original no-noise image.
Best.
K. Frank
1 Like
Hi,
thanks for your answer. I tried MSELoss and also L1Loss, but the images are still grayish. I checked that the result images have around 30-70 different values (instead of 256) all nearing the maximum brightness (getting values from 200-255). I wonder if I should somehow normalize the images during training or something. Do you have any idea what is going wrong?
Br. Veera
Hi Veera!
What sorts of values do you get for MSELoss and L1Loss? Do they
seem consistent with the fact that your “result images” aren’t as good
as you would like?
Both MSELoss and L1Loss should become zero as your result images
becomes exactly the same as your target images. Do you get zero when
you apply your loss function to (result, result) and (target, target)?
Is it possible to overfit your model by training (a lot) on a small subset
of your training data? You would like to be able to get very good “result
images” for just a small set of training samples by training a lot on just
those samples. If you can’t, there is probably something the matter with
your model or your training algorithm.
Normalizing your images – both the inputs and targets – could potentially
improve training. However, unless different samples within your training
set differ a lot in brightness or contrast, or skew towards being unusually
bright, I wouldn’t expect normalization to fix the problem that you are
describing.
Best.
K. Frank