<commit_before>import time
import numpy as np
import logging

logger = logging.getLogger(__file__)


class StereoMRF(object):
    """
    Markov Random Field with loopy belief propagation (min-sum message passing). 
    """
    def __init__(self, dim, n_levels):
        self.n_levels = n_levels
        self.dimension = (n_levels,) + dim
        self._message_field = {'south': np.zeros(self.dimension, dtype=np.float32),
                               'west': np.zeros(self.dimension, dtype=np.float32),
                               'north': np.zeros(self.dimension, dtype=np.float32),
                               'east': np.zeros(self.dimension, dtype=np.float32),
                               'data': np.zeros(self.dimension, dtype=np.float32)}

    def _init_fields(self, image_left, image_right, prior=None, prior_trust_factor=1.0, reinit_messages=True):
        """
        Initialise the message fields -- each hidden variable contains 5 message boxes from the 4 adjacent variables 
        (south, west, north, east) and the observed variable (data).
        
        Args:
            image_left: a numpy array representing the left image (in grayscale with values in [0, 1])
            image_right: a numpy array with the same structure and shape as left image, representing the right image 
            prior: a numpy array with the same shape as left image, providing alternative source of truth (or initial
             belief) about some pixels' corresponding disparity values. 
            prior_trust_factor: float in [0, 1] telling how much the prior should be trusted.
            reinit_messages: whether the messages field should be reset to 0

        Returns:
            In-place method
        """
        assert image_left.shape == image_right.shape
        self.reference_image = image_left.astype('float32')
        self.secondary_image = image_right.astype('float32')
        if reinit_messages:
            self._message_field = {'south': np.zeros(self.dimension, dtype=np.float32),
                                   'west': np.zeros(self.dimension, dtype=np.float32),
                                   'north': np.zeros(self.dimension, dtype=np.float32),
                                   'east': np.zeros(self.dimension, dtype=np.float32),
                                   'data': np.zeros(self.dimension, dtype=np.float32)}

        nrow, ncol = self.reference_image.shape
        # crop the border from the high right in the right image as it cannot be mathed to the left
        if prior is not None:
            assert image_right.shape == prior.shape and prior.shape == self._message_field['data'].shape[1:]
            for l in xrange(self.n_levels):
                data_contrib = np.hstack((np.abs(self.reference_image[:, l:] - self.secondary_image[:, :ncol - l]),
                                          np.ones((nrow, l))))
                # data_contrib /= np.max(data_contrib) * 1/(np.mean(data_contrib))
                # print(data_contrib.min(), data_contrib.max(), data_contrib.mean())
                prior_mask = prior == l

                # equivalent to linear interpolating between the prior regions (having lowest possible value of -1)
                # and the data (weighted by the data_trust_factor = 1 - prior_trust_factor)
                self._message_field['data'][l, prior_mask] = (1 - prior_trust_factor) * data_contrib[prior_mask] \
                                                             - prior_trust_factor
                self._message_field['data'][l, ~prior_mask] = data_contrib[~prior_mask]
        else:
            for l in xrange(self.n_levels):
                self._message_field['data'][l, :, :ncol - l] = np.abs(self.reference_image[:, l:]
                                                                      - self.secondary_image[:, :ncol - l])

    def _update_message_fields(self):
        """
        Pass messages from each hidden and observable variable to the corresponding adjacent hidden. 
        Since messages get summed up, normalise them to prevent overflows.
        
        Returns:
            In-place method
        """
        for direction in ['south', 'west', 'north', 'east']:
            message_updates = np.sum([field for d, field in self._message_field.iteritems() if d != direction], axis=0)
            if direction == 'south':
                self._message_field[direction][:, 1:, :] = message_updates[:, :-1, :]
            elif direction == 'west':
                self._message_field[direction][:, :, :-1] = message_updates[:, :, 1:]
            elif direction == 'north':
                self._message_field[direction][:, :-1, :] = message_updates[:, 1:, :]
            elif direction == 'east':
                self._message_field[direction][:, :, 1:] = message_updates[:, :, :-1]
            # add normalisation to the message values, as they grow exponentially with the number of iterations
            norm_factor = np.max(self._message_field[direction], axis=0, keepdims=True)
            norm_factor[norm_factor == 0] = 1
            self._message_field[direction] /= norm_factor

    def _update_belief_field(self):
        """
        Find an optimal (Maximum A-Posteriori) assignment for each pixel. <commit_msg>fix reference image in observed vars<commit_after>import time
import numpy as np
import logging

logger = logging.getLogger(__file__)


class StereoMRF(object):
    """
    Markov Random Field with loopy belief propagation (min-sum message passing). 
    """
    def __init__(self, dim, n_levels):
        self.n_levels = n_levels
        self.dimension = (n_levels,) + dim
        self._message_field = {'south': np.zeros(self.dimension, dtype=np.float32),
                               'west': np.zeros(self.dimension, dtype=np.float32),
                               'north': np.zeros(self.dimension, dtype=np.float32),
                               'east': np.zeros(self.dimension, dtype=np.float32),
                               'data': np.zeros(self.dimension, dtype=np.float32)}

    def _init_fields(self, image_left, image_right, prior=None, prior_trust_factor=1.0, reinit_messages=True):
        """
        Initialise the message fields -- each hidden variable contains 5 message boxes from the 4 adjacent variables 
        (south, west, north, east) and the observed variable (data).
        
        Args:
            image_left: a numpy array representing the left image (in grayscale with values in [0, 1])
            image_right: a numpy array with the same structure and shape as left image, representing the right image 
            prior: a numpy array with the same shape as left image, providing alternative source of truth (or initial
             belief) about some pixels' corresponding disparity values. 
            prior_trust_factor: float in [0, 1] telling how much the prior should be trusted.
            reinit_messages: whether the messages field should be reset to 0

        Returns:
            In-place method
        """
        assert image_left.shape == image_right.shape
        self.reference_image = image_left.astype('float32')
        self.secondary_image = image_right.astype('float32')
        if reinit_messages:
            self._message_field = {'south': np.zeros(self.dimension, dtype=np.float32),
                                   'west': np.zeros(self.dimension, dtype=np.float32),
                                   'north': np.zeros(self.dimension, dtype=np.float32),
                                   'east': np.zeros(self.dimension, dtype=np.float32),
                                   'data': np.zeros(self.dimension, dtype=np.float32)}

        nrow, ncol = self.reference_image.shape
        # crop the border from the high right in the right image as it cannot be mathed to the left
        if prior is not None:
            assert image_right.shape == prior.shape and prior.shape == self._message_field['data'].shape[1:]
            for l in xrange(self.n_levels):
                data_contrib = np.hstack((np.abs(self.reference_image[:, l:] - self.secondary_image[:, :ncol - l]),
                                          np.ones((nrow, l))))
                # data_contrib /= np.max(data_contrib) * 1/(np.mean(data_contrib))
                # print(data_contrib.min(), data_contrib.max(), data_contrib.mean())
                prior_mask = prior == l

                # equivalent to linear interpolating between the prior regions (having lowest possible value of -1)
                # and the data (weighted by the data_trust_factor = 1 - prior_trust_factor)
                self._message_field['data'][l, prior_mask] = (1 - prior_trust_factor) * data_contrib[prior_mask] \
                                                             - prior_trust_factor
                self._message_field['data'][l, ~prior_mask] = data_contrib[~prior_mask]
        else:
            for l in xrange(self.n_levels):
                self._message_field['data'][l, :, l:] = np.abs(self.reference_image[:, l:]
                                                               - self.secondary_image[:, :ncol - l])

    def _update_message_fields(self):
        """
        Pass messages from each hidden and observable variable to the corresponding adjacent hidden. 
        Since messages get summed up, normalise them to prevent overflows.
        
        Returns:
            In-place method
        """
        for direction in ['south', 'west', 'north', 'east']:
            message_updates = np.sum([field for d, field in self._message_field.iteritems() if d != direction], axis=0)
            if direction == 'south':
                self._message_field[direction][:, 1:, :] = message_updates[:, :-1, :]
            elif direction == 'west':
                self._message_field[direction][:, :, :-1] = message_updates[:, :, 1:]
            elif direction == 'north':
                self._message_field[direction][:, :-1, :] = message_updates[:, 1:, :]
            elif direction == 'east':
                self._message_field[direction][:, :, 1:] = message_updates[:, :, :-1]
            # add normalisation to the message values, as they grow exponentially with the number of iterations
            norm_factor = np.max(self._message_field[direction], axis=0, keepdims=True)
            norm_factor[norm_factor == 0] = 1
            self._message_field[direction] /= norm_factor

    def _update_belief_field(self):
        """
        Find an optimal (Maximum A-Posteriori) assignment for each pixel. 
{'fraction_non_alphanumeric': 0.10462795780670911, 'fraction_numerical': 0.009502993442934524, 'mean_word_length': 2.046902142443544, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 21, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4535059', 'n_tokens_mistral': 2936, 'n_tokens_neox': 2732, 'n_words': 798}