<jupyter_start><jupyter_text>Get a training example<jupyter_code>%%time
image, labels = data[20]

heatmaps = labels['heatmaps']
offsets = labels['offsets']
assert 5 * labels['num_boxes'] == (labels['heatmaps'] == 1).sum()
print('number of persons:', labels['num_boxes'])
masks = labels['masks']
np.unique(masks.numpy())<jupyter_output><empty_output><jupyter_text>Show the image with masks<jupyter_code>image = Image.fromarray((255*image.permute(1, 2, 0).numpy()).astype('uint8'))
width, height = image.size
print(width, height)
assert (width % 32 == 0) and (height % 32 == 0)
image.putalpha(255)
def show_mask(image, mask):

    red = np.array([255, 0, 0], dtype='uint8')
    gray_mask = mask.numpy().astype('uint8')
    color_mask = red * np.expand_dims(gray_mask, 2)

    gray_mask = Image.fromarray(100 * gray_mask)
    color_mask = Image.fromarray(color_mask)
    color_mask.putalpha(gray_mask)
    
    image_copy = image.copy()
    image_copy.alpha_composite(color_mask.resize((width, height)))
    return image_copy
show_mask(image, masks[0])
show_mask(image, masks[1])<jupyter_output><empty_output><jupyter_text>Show the image with heatmaps<jupyter_code>image_copy = image.copy()

colors = {
    0: [255, 0, 0],  # red - top
    1: [0, 0, 255],  # blue - bottom
    2: [255, 255, 0],  # yellow - left
    3: [255, 0, 255],  # pink - right
    4: [0, 255, 0]  # green - center
}

for i, color in colors.items():

    h = np.expand_dims(heatmaps[i].numpy() > 4e-1, 2)
    gray_h = Image.fromarray(255*h[:, :, 0].astype('uint8'))
    color_h = Image.fromarray((color * h).astype('uint8'))
    
    color_h = color_h.resize((width, height))
    color_h.putalpha(gray_h.resize((width, height)))
    
    image_copy.alpha_composite(color_h)
    
image_copy<jupyter_output><empty_output><jupyter_text>Show the image with offsets<jupyter_code>stride = 4

true_peaks = {}
for i in range(5):
    
    h = heatmaps[i].numpy()
    y, x = np.where(h == 1)
    vectors = offsets[(2*i):(2*i + 2)].permute(1, 2, 0).numpy()
    assert (vectors[h < 1.0] == 0).all()
    
    offset_vectors = vectors[y, x]
    vy = offset_vectors[:, 0]
    vx = offset_vectors[:, 1]

    true_y = stride * (y + vy)
    true_x = stride * (x + vx)

    true_peaks[i] = np.stack([true_x, true_y], axis=1)
image_copy = image.copy()
draw = ImageDraw.Draw(image_copy)
s = 3

for i, peaks in true_peaks.items():
    for x, y in peaks:
        color = tuple(colors[i])
        draw.ellipse([(x - s, y - s), (x + s, y + s)], fill=color)

image_copy<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16837136113296616, 'fraction_numerical': 0.03619197482297404, 'mean_word_length': 3.966796875, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '411355', 'license': 'MIT', 'n_tokens_mistral': 1054, 'n_tokens_neox': 979, 'n_words': 206, 'path': 'data/test_input_pipeline.ipynb', 'repo_name': 'TropComplique/ExtremeNet', 'size': '2542'}