<jupyter_start><jupyter_text>Generate a training set of images and annotationsWe'll use the [PASCAL VOC2012](http://host.robots.ox.ac.uk/pascal/VOC/voc2012/) and [PASCAL-Part](http://www.stat.ucla.edu/~xianjie.chen/pascal_part_dataset/pascal_part.html) images and annotations as a training set for UNet.Let's download both datasets first:<jupyter_code># Uncomment the following lines to download and extract VOC2012 and PASCAL-Part data
# !mkdir data
# !wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar -P ./data
# !tar xf ./data/VOCtrainval_11-May-2012.tar -C ./data
# !rm -rf ./data/VOCtrainval_11-May-2012.tar
# !wget http://www.stat.ucla.edu/~xianjie.chen/pascal_part_dataset/trainval.tar.gz -P ./data
# !tar xfz ./data/trainval.tar.gz -C ./data
# !rm -rf ./data/trainval.tar.gz
import numpy as np
import os
import pickle
import warnings

from matplotlib.pyplot import imshow
from skimage.color import rgb2gray
from skimage.io import imread
from skimage.transform import resize

from anno import ImageAnnotation

warnings.filterwarnings('ignore')
%matplotlib inline<jupyter_output><empty_output><jupyter_text>Structure of input and output tensorsWe'll feed the Convolutional Neural Network with 368x368 normalized grayscale frames (pixel values between 0 and 1). UNet seems to work better with a single input channel, but RGB input should be explored too. The output tensor consists of two 368x368 binary matrices. On the first one, a 1 means that the object to be segmented is being detected on that pixel or 0 otherwise. The second matrix is the opposite of the first, a 1 means that background is present in that pixel and a 0 means that the object is there.We'll generate ground-truth data of cars to train a car-segmentation UNet model. Parse VOC2012 annotationsFirst we get a list of images that are segmented.<jupyter_code>original_dir = './data/VOCdevkit/VOC2012/JPEGImages'
segmented_dir = './data/VOCdevkit/VOC2012/SegmentationClass'
with open('./data/VOCdevkit/VOC2012/ImageSets/Segmentation/trainval.txt', 'r') as f:
    segmented = [line.strip() for line in f.readlines()]<jupyter_output><empty_output><jupyter_text>Let's analyze a segmented car image:<jupyter_code>name = '2007_001321'
original_img = imread(os.path.join(original_dir, '{}.jpg'.format(name)))
segment_img = imread(os.path.join(segmented_dir, '{}.png'.format(name)))
imshow(original_img)
imshow(segment_img)
colors = set()
for i in range(segment_img.shape[0]):
    for j in range(segment_img.shape[0]):
        r, g, b = segment_img[i, j, :]
        colors.add((r, g, b))
print(colors)<jupyter_output>{(224, 224, 192), (0, 0, 0), (128, 128, 128)}<jupyter_text>Cars are segmented with gray (128, 128, 128) pixels in the VOC2012 annotations. We can then generate the correct output matrix for this segment pretty easily:<jupyter_code>mask = np.all(segment_img == (128, 128, 128), axis=-1)
output_segment = np.zeros(mask.shape)
output_segment[mask] = 1.
imshow(output_segment)<jupyter_output><empty_output><jupyter_text>Let's also explore a handy calculation which returns the percentage of the image covered by a car:<jupyter_code>output_segment.sum() / output_segment.shape[0] / output_segment.shape[1]<jupyter_output><empty_output><jupyter_text>56% of the image is covered by the car, this is useful when cars only fill a small percentage of the image and are too small to be segmented.The last preprocess needed is converting the input to grayscale and resizing both input and output to a 368x368xChannels shape.<jupyter_code>height = 368
width = 368
input_img = resize(rgb2gray(original_img), (height, width), anti_aliasing=True).reshape((height, width, 1))
imshow(input_img[:, :, 0])
input_img.shape
output_segment = resize(output_segment, (height, width), anti_aliasing=True)
output_segment[output_segment > 0.] = 1.
output_back = np.ones(output_segment.shape)
output_back[output_segment > 0.] = 0.
output = np.stack([output_segment, output_back], axis=-1)
imshow(output[:,:,0])
output.shape<jupyter_output><empty_output><jupyter_text>Generate VOC2012 train setNow it's time to iterate through all the segmented images and add the ones with large enough cars to our train set.Let's define the general constants to generate the train set which could be modified to generate a set of segments of any other class.<jupyter_code># Input/output sizes
height = 368
width = 368
# Class we want to detect
segment_color = (128, 128, 128)  # Cars
# Percentage of the image covered by this class
threshold = 0.05
input_images = []
output_tensors = []
for name in segmented:
    original_img = imread(os.path.join(original_dir, '{}.jpg'.format(name)))
    segment_img = imread(os.path.join(segmented_dir, '{}.png'.format(name)))
    mask = np.all(segment_img == segment_color, axis=-1)
    output_segment = np.zeros(mask.shape)
    output_segment[mask] = 1.
    if output_segment.sum() / output_segment.shape[0] / output_segment.shape[1] > threshold:
        image = resize(rgb2gray(original_img), (height, width), anti_aliasing=True).reshape((1, height, width, 1))
        output_segment = resize(output_segment, (height, width), anti_aliasing=True)
        output_segment[output_segment > 0.] = 1.
        output_back = np.ones(output_segment.shape)
        output_back[output_segment > 0.] = 0.
        tensor = np.stack([output_segment, output_back], axis=-1).reshape((1, height, width, 2))
        input_images.append(image)
        output_tensors.append(tensor)
print("Images collected: {}".format(len(input_images)))<jupyter_output>Images collected: 147<jupyter_text>Let's check we have gathered the correct images:<jupyter_code>imshow(input_images[32][0, :, :, 0])
imshow(output_tensors[32][0, :, :, 0])<jupyter_output><empty_output><jupyter_text>Parse PASCAL-Part DatasetThe 147 images from VOC2012 may be not sufficient to train a good model. There's another useful annotated dataset that we've downloaded: [PASCAL-Part](http://www.stat.ucla.edu/~xianjie.chen/pascal_part_dataset/pascal_part.html). It is based on the VOC2010 data, so we expect some of the images to be repeated. We only save those that aren't already part of the VOC2012 segmented images.This dataset comes in `.mat` format, we use slightly modified [pascal-part-py](https://github.com/twuilliam/pascal-part-py) helpers to quickly parse these files.<jupyter_code>annotations_dir = './data/Annotations_Part/'
voc2012 = set(segmented)
segmented = []
for annotation in os.listdir(annotations_dir):
    name = annotation.split('.')[0]
    if name in voc2012:
        continue
    ann = ImageAnnotation(
        os.path.join(original_dir, '{}.jpg'.format(name)),
        os.path.join(annotations_dir, '{}.mat'.format(name))
    )
    if 'car' in [ann.objects[i].class_name for i in range(ann.n_objects)]:
        segmented.append(name)
print("There is an additional number of {0} images of cars in the PASCAL-Part dataset.".format(len(segmented)))<jupyter_output>There is an additional number of 877 images of cars in the PASCAL-Part dataset.<jupyter_text>Generate PASCAL-Part train setLet's add those images where cars occupy a large enough portion of the image to the train set:<jupyter_code>for name in segmented:
    ann = ImageAnnotation(
        os.path.join(original_dir, '{}.jpg'.format(name)),
        os.path.join(annotations_dir, '{}.mat'.format(name))
    )
    if ann.part_mask.sum() / ann.part_mask.shape[0] / ann.part_mask.shape[1] > threshold:
        image = resize(rgb2gray(ann.im), (height, width), anti_aliasing=True).reshape(1, height, width, 1)
        output_segment = resize(ann.part_mask, (height, width), anti_aliasing=True)
        output_segment[output_segment > 0.] = 1.
        output_back = np.ones(output_segment.shape)
        output_back[output_segment > 0.] = 0.
        tensor = np.stack([output_segment, output_back], axis=-1).reshape((1, height, width, 2))
        input_images.append(image)
        output_tensors.append(tensor)<jupyter_output><empty_output><jupyter_text>Again, let's check we have gathered the correct images:<jupyter_code>imshow(input_images[-1][0, :, :, 0])
imshow(output_tensors[-1][0, :, :, 0])
print("Images collected: {}".format(len(input_images)))<jupyter_output>Images collected: 665<jupyter_text>Concatenate and pickle final datasetLet's save input and output tensors as *ndarray*s of shape (N, height, width, channels) and pickle them for training UNet networks.<jupyter_code>input_images = np.concatenate(input_images, axis=0)
output_tensors = np.concatenate(output_tensors, axis=0)
print("Input tensor shape: {}".format(input_images.shape))
print("Output tensor shape: {}".format(output_tensors.shape))
with open('train_images.pkl', 'wb') as fout:
    pickle.dump(input_images, fout, protocol=4)
with open('train_labels.pkl', 'wb') as fout:
    pickle.dump(output_tensors, fout, protocol=4)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11740160144355476, 'fraction_numerical': 0.02830720649599639, 'mean_word_length': 5.115862068965518, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 53, 'https://': 1, 'lorem ipsum': 0, 'www.': 3, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '32865', 'license': 'Apache-2.0', 'n_tokens_mistral': 3085, 'n_tokens_neox': 2840, 'n_words': 808, 'path': 'train-data-generator.ipynb', 'repo_name': 'gmontamat/semantic-segmentation-unet', 'size': '8867'}