media/deep-learning/semantic-segmentation/make-labels.py.Semantic Segmentation with U-Net
deep-learning
convolutional-neural-networks
computer-vision
semantic-segmentation
u-net
transpose-convolution
Labeling every pixel rather than drawing a box, the transpose convolution that grows a representation back, and the U-Net that puts the two together.
Three tasks sit on a ladder. Object recognition answers what is in a picture, such as whether it shows a cat. Object detection goes a step further and reports objects with rectangles, which is what Bounding Box Predictions with YOLO does, producing a box, a class, and a confidence for each object it finds. For a great many applications that is enough.
A box is a coarse answer, though. It says a car is somewhere inside this rectangle, and the rectangle inevitably contains road, sky, and pieces of other cars. Semantic segmentation asks the harder question, which is not where the object roughly is but which exact pixels belong to it. The answer is a label for every pixel in the image.
Pixel Labels
Start with the simplest version of the task, separating a car from everything else. Every pixel gets one of two labels, 1 if it is part of the car and 0 if it is not. The output is therefore not a class, and not four coordinates, but a grid of labels the same height and width as the input.
A richer version uses four labels. Keep 0 for anything outside the named scene classes, label the car 1, the buildings 2, and the road 3. The algorithm then has to color in the whole picture, deciding which label every pixel carries.
Those labels are easier to believe once they are drawn on the picture itself. Coarsen a scene into a grid of blocks, give every block the label its pixels mostly carry, and each block of the photograph ends up with a digit sitting on it. Lift the digits off the image and what is left standing on its own is the segmentation map, which is the thing the network is trained to produce.
The top row is the two label version, where a block is 1 if it belongs to the car and 0 if it does not. The bottom row is the richer version. Label 0 covers everything outside the three named scene classes, while the car, building, and road get labels of their own.
Blocks are used here only so that the digits stay large enough to read, since a real segmentation map carries one label per pixel rather than one per block. Set against what a detector has to say about a street, the difference in what is being asked for is immediate.
media/deep-learning/semantic-segmentation/make-labels.py.The middle panel is everything a detector has to say about this street. Seven rectangles, each holding a vehicle, and each also holding a strip of road, a slice of sky, and pieces of whatever is parked behind. The right panel is everything a segmentation network has to say about it, which is one label at each of the 921 600 pixels, so the boundary of a car follows the car and the road is marked wherever it is genuinely visible between them.
Those labels are predictions rather than human annotation, and the difference shows. The queue of vehicles near the center of the frame merges into a single region instead of separating into individual cars, because semantic segmentation labels a pixel with its class and has no notion of which car a pixel belongs to.
Formally, for an image of height \(h\) and width \(w\) the output has a prediction at all \(h \times w\) positions. With \(n_{\text{classes}}\) classes, the final layer produces a vector of \(n_{\text{classes}}\) numbers at each pixel, and taking the largest entry of that vector assigns the pixel to a class. An output of shape \(h \times w \times 3\) is therefore three numbers per pixel, not three pixels.
This is worth doing because several applications need exactly it, and a box would be useless for them.
A self driving car needs to know which pixels are drivable surface. Drawing a rectangle around the road is meaningless, since a road is not rectangular and the rectangle would include the pavement and the cars parked on it. Marking every pixel that can be driven over is the useful answer, and some self driving teams use segmentation for precisely that.
In medical imaging the same shift matters more. Given a chest X-ray, a diagnosis is useful, but an outline of exactly which pixels are lung, which are heart, and which are collarbone is more useful still, because it makes irregularities easier to spot and helps surgeons plan. Novikov et al. (2018) do exactly that, giving each structure its own color. Given a brain MRI, outlining a tumor by hand is slow and tedious work, and an algorithm that segments it automatically saves a radiologist a great deal of time and gives a surgical team something to plan against. That is the result reported by Dong et al. (2017), and the architecture behind it is the one built up in the rest of this page.
The shape of the network that produces such a map follows from the shape of the map itself. A recognition network passes an image forward through layer after layer, and the height and width of its activations shrink the whole way, until nothing is left but a vector of class scores. Taking the last few layers off removes that vector, and what remains is still far smaller than the input. The second half of the network therefore has to undo the shrinking, growing the height and width back up until there is a position for every pixel of the original image.
Everything in the first half is an ordinary ConvNet, of the kind the earlier pages built. The one operation in the second half that has not been covered yet is the one that makes a small grid of activations bigger.
Review Questions
1. How is semantic segmentation more detailed than object detection?
TipAnswer
Object detection identifies an object and places a rectangle around it, so the answer is a handful of numbers per object. Semantic segmentation labels every pixel, so it can trace the exact boundary of an object, a road, or a building. A box says roughly where something is, while a segmentation says which pixels it is made of.
1. Why can the output of an ordinary classifier not serve as a segmentation map?
TipAnswer
Because of its shape. A classifier reduces an image to a short vector, one number per class, and along the way the height and width of its activations shrink steadily. A segmentation map needs a prediction at every pixel position, so it has to be as tall and as wide as the input image. Getting there means the network has to grow the representation back up after shrinking it, which no ordinary classifier does.
1. A segmentation network for three classes outputs a volume of shape \(h \times w \times 3\). What sits at one position of that volume, and how does it become a label?
TipAnswer
Three numbers, one per class, describing how strongly that pixel belongs to each. The label is the position of the largest of the three, which is an \(\arg\max\) over the class axis. Doing that at every one of the \(h \times w\) positions turns the volume into a single map of labels, which is what gets displayed as a colored overlay.
1. Which of these is semantic segmentation? (a) Locating an object of a certain class by drawing a bounding box around it. (b) Locating objects by predicting, for every pixel, which class it belongs to. (c) Locating objects of different classes by drawing bounding boxes around them.
TipAnswer
(b). The answer is a class label at every pixel, so the shape of an object is traced rather than enclosed. (a) is localization and (c) is detection, and both stop at rectangles.
Transpose Convolutions
An ordinary convolution shrinks a representation, or at best keeps it the same size. Segmentation needs the opposite at some point, an operation that takes a small grid of activations and produces a larger one. The transpose convolution is that operation.
Its mechanics invert the usual convolution in one specific way. A normal convolution places the filter on the input, multiplies element by element, and sums to produce one output number. A transpose convolution places the filter on the output. It takes one input number, multiplies the entire filter by that single number, and writes the scaled filter into the output grid.
Because the input positions step across the output by the stride, and the filter is wider than the stride, neighboring placements land on overlapping regions. Where placements overlap, their values are added.
Carrying that through with actual numbers is the quickest way to be sure of it. Take a 2 by 2 input, a 3 by 3 filter, a stride of 2, and a padding of 1, and ask for a 4 by 4 output. The padding applies to the output rather than the input here, so the canvas being written into is 6 by 6, and the one cell border around it is discarded at the end.
The figure below is interactive. Press Play and the four scaled filters are written onto the canvas one at a time, so the totals in the cells grow as the blocks land rather than sitting there finished, and clicking an entry of the 2 by 2 input jumps straight to the step where that entry is placed. Hovering a cell of the canvas works backwards instead. It splits that cell into the contributions that landed on it, which is where the addition becomes visible, and lights up the input entry and the filter weight behind each contribution in the same color, so a result can be traced back to what produced it. Hovering an entry of the filter goes the other way, following that one weight forward to the single cell it writes into.
Follow one number through. The input entry in the top left is 2, so the whole filter is multiplied by 2, giving \([[2, 4, 2], [4, 0, 2], [0, 4, 2]]\), and that 3 by 3 block is written into the top left corner of the canvas. Its first row and its first column fall in the padding, so five of its nine values are thrown away and only four survive. The next input entry along is 1, so the filter is multiplied by 1 and written two columns further across, because the stride is 2. Since the block is 3 wide and the step is only 2, its left column lands on top of the previous block’s right column, and those two values are added rather than one replacing the other. Repeating for the remaining two entries fills the rest of the canvas.
The 4 by 4 left inside the dashed rectangle is the answer.
\[ \begin{bmatrix} 0 & 4 & 0 & 1 \\ 10 & 7 & 6 & 3 \\ 0 & 7 & 0 & 2 \\ 6 & 3 & 4 & 2 \end{bmatrix} \]
Seven of the sixteen output cells receive more than one contribution, and they form a cross through the second row and the second column, which is where the four placements meet. The busiest is the cell holding 7, which takes 2 from the block scaled by 2, 0 from the block scaled by 1, 3 from the block scaled by 3, and 2 from the block scaled by 2 again, and \(2 + 0 + 3 + 2 = 7\).
NoteWhy a library can return a different 4 by 4
The padding above is stripped symmetrically, one cell from every side of the 6 by 6 canvas. A framework does not have to place it that way. Passing this same input, filter, and stride to tf.nn.conv2d_transpose with padding='SAME' and a 4 by 4 output shape returns the top left 4 by 4 of the same canvas, which is
\[ \begin{bmatrix} 2 & 4 & 3 & 2 \\ 4 & 0 & 4 & 0 \\ 3 & 10 & 7 & 6 \\ 6 & 0 & 7 & 0 \end{bmatrix} \]
rather than the grid inside the dashed rectangle. The total padding TensorFlow needs here is a single cell, and it puts that cell on the bottom and the right instead of splitting it between the two sides. The canvas is identical, the operation is identical, and only the choice of which 4 by 4 to keep differs.
One question is worth answering before moving on, which is why bother with this rather than simply scaling the image up. The filter values here are learned, like any other convolutional filter. A fixed rule for enlargement, such as repeating each pixel or interpolating between them, treats every region the same way. A learned filter can decide how a compact representation should be spread into a larger one, which is what makes the operation useful inside a network that has to reconstruct fine detail.
Review Questions
1. What happens when two scaled copies of the filter write to the same output position?
TipAnswer
The contributions are added. A transpose convolution does not choose between placements or overwrite one with another, it accumulates every contribution landing on a position. In the worked example seven of the sixteen output cells receive more than one copy, and the cell holding 7 receives four of them, \(2 + 0 + 3 + 2\).
1. Why does a stride of 2 with a 3 by 3 filter guarantee overlap?
TipAnswer
Because the filter is wider than the step. Each placement covers three columns of the output, while consecutive placements start only two columns apart, so the third column of one placement lands on the first column of the next. Overlap disappears only when the stride is at least as large as the filter, which is also when the operation stops blending neighboring contributions.
1. A transpose convolution has padding 1 and stride 2. The input is \(\begin{bmatrix} 1 & 3 \\ 2 & 4 \end{bmatrix}\) and the filter is \(\begin{bmatrix} 1 & 0 & 1 \\ 0 & 0 & 0 \\ 1 & 0 & 1 \end{bmatrix}\). Fill in \(X\), \(Y\), and \(Z\) in the 4 by 4 output.
\[ \begin{bmatrix} 0 & 0 & 0 & 0 \\ 0 & X & 0 & 7 \\ 0 & 0 & 0 & Y \\ 0 & Z & 0 & 4 \end{bmatrix} \]
TipAnswer
\(X = 10\), \(Y = 0\), \(Z = 6\).
This filter has ones only at its four corners and zeros everywhere else, so each scaled copy writes to four positions and leaves the rest alone. The four copies land two cells apart on a 6 by 6 canvas, and the one cell border is discarded at the end.
\(X\) sits at the only position all four copies reach, so it collects a corner from each, \(1 + 3 + 2 + 4 = 10\). \(Y\) is reached by one copy only, and by a zero inside that copy, so it stays 0. \(Z\) is reached by two copies, taking \(2 \times 1\) from the one scaled by 2 and \(4 \times 1\) from the one scaled by 4, which is 6.
The interactive figure above computes exactly this, with a different filter.
1. Why is a learned transpose convolution preferable to simply enlarging the activations with a fixed rule?
TipAnswer
Because a fixed rule cannot adapt. Repeating pixels or interpolating between them applies the same recipe everywhere, regardless of what the activations mean. The filter in a transpose convolution is trained like any other, so the network can learn how a compact representation should be spread out to reconstruct the detail that matters for the task, which for segmentation means the boundary between one class and the next.
Compression, Expansion, and Skip Connections
With the enlarging operation in hand, the shape of the architecture follows from what the task needs.
The first half of the network is an ordinary ConvNet with its last few layers taken off. A recognition network ends by collapsing its activations into a single vector of class scores, and that ending is the one part a segmentation network has no use for. Everything before it is kept. Convolutions and pooling layers steadily reduce the height and width of the activations while the number of channels grows, so the representation becomes small and deep. That is exactly what recognizing content requires, and it is also lossy in a specific way. A small deep volume can carry the information that something car like is present in the left half of the image, without carrying where the edge of the car runs.
The second half undoes the shrinking. Transpose convolutions grow the height and width back up, step by step, and the number of channels comes back down as they do, so the descent is retraced in reverse until the activations are as large as the input image and a prediction can be made at every pixel.
An architecture built from those two halves alone would work poorly, and the reason is the loss just described. By the time the representation reaches the bottom, the fine spatial detail is gone, so the expanding half is asked to reconstruct a precise boundary from information that no longer contains it.
The fix is the skip connection, and it is the reason the architecture is worth a name. Activations from a layer on the contracting side are copied directly across to the matching layer on the expanding side, skipping everything in between. The later layer then has two sources of information at once.
The high level context arrives up through the expanding path. It is low resolution but semantically rich, and it knows there is a car in this region.
The fine detail arrives across the skip connection. It is high resolution and semantically thin, and it knows how much bodywork like texture sits at this exact pixel, and exactly where it stops and the road begins.
Deciding whether one particular pixel belongs to the car needs both, which is why the two are concatenated and handed to the convolutions that follow.
Review Questions
1. What is lost as pooling makes the activation grid smaller, and why does that matter for segmentation?
TipAnswer
Precise spatial detail. A small deep volume can represent that a car is roughly in one region, but it cannot represent where the boundary of the car falls, because there are no longer enough positions to say so. For classification that loss costs nothing, since the answer is one label for the whole image. For segmentation it is fatal on its own, because the answer needs a decision at every pixel, and that is what the skip connections restore.
1. What two kinds of information meet at a skip connection, and why is neither enough alone?
TipAnswer
Low resolution context from the expanding path, which knows what is in the region but not exactly where, and high resolution detail copied from the contracting path, which knows the texture at a pixel but not what object it belongs to. Deciding whether one pixel belongs to the car needs both at once, so the two are combined and passed through further convolutions. Context alone gives blurry blobs, and detail alone gives edges with no idea what they enclose.
1. The input to a U-Net is \(h \times w \times 3\), the 3 being the red, green, and blue channels. What shape is the output?
TipAnswer
\(h \times w \times n_{\text{classes}}\).
The height and width come back to the size of the input, because a prediction is needed at every pixel of the original image. The depth is set by the final \(1 \times 1\) convolution, which maps whatever channels arrive into one score per class. It has nothing to do with the number of input channels or the number of filters used along the way.
Full Architecture
Putting all of it together gives the shape the architecture is named after. Drawn with the contracting path descending on the left and the expanding path rising on the right, the diagram looks like the letter U.
Reading it from the input, the input is an image of shape \(h \times w \times 3\), the three channels being red, green, and blue. Convolutions, each followed by a ReLU, process it while max pooling occasionally halves the height and width. The channel count grows as the spatial size shrinks, so the volume gets thinner and deeper as it descends, until the bottom of the U holds a small, deep, context rich representation.
Coming back up, each step is a transpose convolution that enlarges the height and width while thinning the volume back down, followed by a skip connection that copies the matching activations from the left side and concatenates them, followed by ordinary convolutions with ReLU that mix the two sources together. That triple repeats until the activations are back to \(h \times w\), where a couple more ordinary convolutions run at full resolution.
The last operation is a \(1 \times 1\) convolution, which looks at one pixel position at a time across all channels and maps it to \(n_{\text{classes}}\) numbers. The output is therefore \(h \times w \times n_{\text{classes}}\), a vector of class scores at every pixel, and an \(\arg\max\) over that vector at each position produces the segmentation map. For three classes the output is \(h \times w \times 3\), and for ten classes it is \(h \times w \times 10\), with the height and width unchanged in both cases.
The architecture is due to Ronneberger, Fischer, and Brox (2015), who were working on biomedical image segmentation at the time. The ideas turned out to transfer well beyond that, and U-Net is now one of the standard architectures for segmentation of any kind.
Review Questions
1. Why do the activations become deeper as they become smaller on the contracting path?
TipAnswer
Because the two changes come from different layers. Pooling reduces the height and width, while each convolutional layer can use more filters than the last, which increases the channel count. The combination trades spatial resolution for representational depth, giving a compact volume that describes what is present rather than exactly where it is. The expanding path runs the trade the other way, growing the height and width while the channel count falls.
1. What does the final \(1 \times 1\) convolution do, and why is that the right last step?
TipAnswer
It maps the channels at each pixel position to one score per class, turning an \(h \times w \times \text{channels}\) volume into an \(h \times w \times n_{\text{classes}}\) volume. A \(1 \times 1\) filter looks at a single position across all channels and mixes them without touching neighboring pixels, which is exactly what is wanted for a per-pixel decision at this point, since the spatial mixing has already been done by the convolutions before it.
1. Why is the network drawn as a U rather than as a straight line of layers?
TipAnswer
Because the drawing puts the shrinking half on the way down and the growing half on the way up, so the spatial size of the activations traces a U, small at the bottom and full size at both ends. Drawing it that way also puts each skip connection on a horizontal line between two layers of equal size, which is what makes the copying across visible in the picture.
References
- Dong, H., Yang, G., Liu, F., Mo, Y., & Guo, Y. (2017). Automatic brain tumor detection and segmentation using U-Net based fully convolutional networks. In Medical Image Understanding and Analysis (pp. 506-517). Springer. https://doi.org/10.1007/978-3-319-60964-5_44
- Novikov, A. A., Lenis, D., Major, D., Hladuvka, J., Wimmer, M., & Buhler, K. (2018). Fully convolutional architectures for multiclass segmentation in chest radiographs. IEEE Transactions on Medical Imaging, 37(8), 1865-1876. https://doi.org/10.1109/TMI.2018.2806086
- Ronneberger, O., Fischer, P., & Brox, T. (2015). U-Net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention, MICCAI 2015, Lecture Notes in Computer Science (Vol. 9351, pp. 234-241). Springer. https://doi.org/10.1007/978-3-319-24574-4_28