<commit_before>from __future__ import print_function

import errno
import gzip
import math
import os

import numpy as np
import skimage.color

from scipy import ndimage
from six.moves import zip_longest
from six.moves import cPickle as pickle

def split(image):
    """Split the image data into the top and bottom half."""
    split_height = image.shape[0] / 2
    return image[:split_height], image[split_height:]

BYTE_MAX = 255
CHANNEL_MAX = np.float32(8)
MAX_RED_VALUE = BYTE_MAX - CHANNEL_MAX
CHANNELS_MAX = CHANNEL_MAX * CHANNEL_MAX
MAX_DEPTH = MAX_RED_VALUE * CHANNELS_MAX
COLOR_CHANNELS = 3

def decode_depth(image):
    """~14 bits of depth in millimeters is encoded with 8 bits in red and 3 bits in each of green and blue."""
    orientation = [1, 0, 0, 0] # default orientation if not present in image.
   
    if np.array_equal(image[0, 0], [BYTE_MAX, 0, 0, BYTE_MAX]):
        # Orientation quaternion is present.
        pixel = image[0, 1]
        for c in range(len(orientation)):
            orientation[c] = ((2.0 * pixel[c]) / BYTE_MAX) - 1

        # Clear out the pixels so they don't get interepreted as depth.
        image[0, 0] = [0, 0, 0, BYTE_MAX]
        image[0, 1] = [0, 0, 0, BYTE_MAX]

    red = image[:, :, 0]
    green = image[:, :, 1]
    blue = image[:, :, 2]

    depth = ((MAX_RED_VALUE - red) * CHANNELS_MAX) + ((green - red) * CHANNEL_MAX) + (blue - red)
    
    # Zero in the red channel indicates the sensor provided no data.
    depth[np.where(red == 0)] = float('nan')
    return depth, orientation

def load_image(image_path):
    combined_image = ndimage.imread(image_path).astype(np.float32)
    color_image, depth_image = split(combined_image)
    color_image = color_image[:, :, 0 : COLOR_CHANNELS] / BYTE_MAX # Discard alpha and normalize
    depths, attitude = decode_depth(depth_image)
    return (color_image, depths, attitude)

def ascending_factors(number):
    factor = 2
    while number > 1:
        if number % factor == 0:
            yield factor
            number = number / factor
        else:
            factor += 1

def compute_scales(height, width):
    height_scales = reversed(list(ascending_factors(height)))
    width_scales = reversed(list(ascending_factors(width)))
    return list(zip_longest(height_scales, width_scales, fillvalue=1))

def mipmap_imputer(image, strategy=np.mean, scales=None):
    """Fill NaNs with localized aggregate values using mipmaps"""
    # Combination of: http://stackoverflow.com/questions/14549696/mipmap-of-image-in-numpy
    # and: http://stackoverflow.com/questions/5480694/numpy-calculate-averages-with-nans-removed
    scales = scales if scales else compute_scales(image.shape[0], image.shape[1])<commit_msg>Add function to encode depth back into a pixel.<commit_after>from __future__ import print_function

import errno
import gzip
import math
import os

import numpy as np
import skimage.color

from scipy import ndimage
from six.moves import zip_longest
from six.moves import cPickle as pickle

def split(image):
    """Split the image data into the top and bottom half."""
    split_height = image.shape[0] / 2
    return image[:split_height], image[split_height:]

BYTE_MAX = 255
CHANNEL_MAX = np.float32(8)
MAX_RED_VALUE = BYTE_MAX - CHANNEL_MAX
CHANNELS_MAX = CHANNEL_MAX * CHANNEL_MAX
MAX_DEPTH = MAX_RED_VALUE * CHANNELS_MAX
COLOR_CHANNELS = 3

def decode_depth(image):
    """~14 bits of depth in millimeters is encoded with 8 bits in red and 3 bits in each of green and blue."""
    orientation = [1, 0, 0, 0] # default orientation if not present in image.
   
    if np.array_equal(image[0, 0], [BYTE_MAX, 0, 0, BYTE_MAX]):
        # Orientation quaternion is present.
        pixel = image[0, 1]
        for c in range(len(orientation)):
            orientation[c] = ((2.0 * pixel[c]) / BYTE_MAX) - 1

        # Clear out the pixels so they don't get interepreted as depth.
        image[0, 0] = [0, 0, 0, BYTE_MAX]
        image[0, 1] = [0, 0, 0, BYTE_MAX]

    red = image[:, :, 0]
    green = image[:, :, 1]
    blue = image[:, :, 2]

    depth = ((MAX_RED_VALUE - red) * CHANNELS_MAX) + ((green - red) * CHANNEL_MAX) + (blue - red)
    
    # Zero in the red channel indicates the sensor provided no data.
    depth[np.where(red == 0)] = float('nan')
    return depth, orientation

def encode_normalized_depth(depth):
    depth_in_mm = int(depth * MAX_DEPTH)
    red_bits = depth_in_mm / int(CHANNELS_MAX)
    lower_bits = depth_in_mm % int(CHANNELS_MAX)
    green_bits = lower_bits / int(CHANNEL_MAX)
    blue_bits = lower_bits % int(CHANNEL_MAX)
    red = int(MAX_RED_VALUE) - red_bits
    return [red, red + green_bits, red + blue_bits, BYTE_MAX]

def load_image(image_path):
    combined_image = ndimage.imread(image_path).astype(np.float32)
    color_image, depth_image = split(combined_image)
    color_image = color_image[:, :, 0 : COLOR_CHANNELS] / BYTE_MAX # Discard alpha and normalize
    depths, attitude = decode_depth(depth_image)
    return (color_image, depths, attitude)

def ascending_factors(number):
    factor = 2
    while number > 1:
        if number % factor == 0:
            yield factor
            number = number / factor
        else:
            factor += 1

def compute_scales(height, width):
    height_scales = reversed(list(ascending_factors(height)))
    width_scales = reversed(list(ascending_factors(width)))
    return list(zip_longest(height_scales, width_scales, fillvalue=1))

def mipmap_imputer(image, strategy=np.mean, scales=None):
    """Fill NaNs with localized aggregate values using mipmaps"""
    # Combination of: http://stackoverflow.com/questions/14549696/mipmap-of-image-in-numpy
    # and: http://stackoverflow.com/questions/5480694/numpy-calculate-averages-with-nans-removed
    scales = scales if scales else compute_scales(image.shape[0], image.shape[1])
{'fraction_non_alphanumeric': 0.10700222564629344, 'fraction_numerical': 0.022256462934429035, 'mean_word_length': 3.4903920061491163, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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}}
{'dir': 'git-commits-cleaned', 'id': '6130142', 'n_tokens_mistral': 2040, 'n_tokens_neox': 1898, 'n_words': 552}