/ ****************************************************************************/

# -*- coding: utf-8 -*-
# This file is part of the Horus Project

__author__ = 'Jes煤s Arroyo Torrens <jesus.arroyo@bq.com>'
__copyright__ = 'Copyright (C) 2014-2016 Mundo Reader S.L.'
__license__ = 'GNU General Public License v2 http://www.gnu.org/licenses/gpl2.html'

import cv2
import math
import numpy as np
import scipy.ndimage

from horus import Singleton
from horus.engine.calibration.calibration_data import CalibrationData
from horus.engine.algorithms.point_cloud_roi import PointCloudROI


@Singleton
class LaserSegmentation(object):

    def __init__(self):
        self.calibration_data = CalibrationData()
        self.point_cloud_roi = PointCloudROI()

        self.red_channel = 'R (RGB)'
        self.threshold_enable = False
        self.threshold_value = 0
        self.blur_enable = False
        self.blur_value = 0
        self.window_enable = False
        self.window_value = 0
        self.refinement_method = 'SGF'

    def set_red_channel(self, value):
        self.red_channel = value

    def set_threshold_enable(self, value):
        self.threshold_enable = value

    def set_threshold_value(self, value):
        self.threshold_value = value

    def set_blur_enable(self, value):
        self.blur_enable = value

    def set_blur_value(self, value):
        self.blur_value = 2 * value + 1

    def set_window_enable(self, value):
        self.window_enable = value

    def set_window_value(self, value):
        self.window_value = value

    def set_refinement_method(self, value):
        self.refinement_method = value

    def compute_2d_points(self, image):
        if image is not None:
            image = self.compute_line_segmentation(image)
            # Peak detection: center of mass
            s = image.sum(axis=1)
            v = np.where(s > 0)[0]
            u = (self.calibration_data.weight_matrix * image).sum(axis=1)[v] / s[v]
            if self.refinement_method == 'SGF':
                # Segmented gaussian filter
                u, v = self._sgf(u, v, s)
            elif self.refinement_method == 'RANSAC':
                # Random sample consensus
                u, v = self._ransac(u, v)
            return (u, v), image

    def compute_hough_lines(self, image):
        if image is not None:
            image = self.compute_line_segmentation(image)
            lines = cv2.HoughLines(image, 1, np.pi / 180, 120)
            # if lines is not None:
            #   rho, theta = lines[0][0]
            #   ## Calculate coordinates
            #   u1 = rho / np.cos(theta)
            #   u2 = u1 - height * np.tan(theta)
            return lines

    def compute_line_segmentation(self, image, roi_mask=False):
        if image is not None:
            # Apply ROI mask
            if roi_mask:
                image = self.point_cloud_roi.mask_image(image)
            # Obtain red channel
            image = self._obtain_red_channel(image)
            if image is not None:
                # Threshold image
                image = self._threshold_image(image)
                # Window mask
                image = self._window_mask(image)
            return image

    def _obtain_red_channel(self, image):
        ret = None
        if self.red_channel == 'R (RGB)':
            ret = cv2.split(image)[0]
        elif self.red_channel == 'Cr (YCrCb)':
            ret = cv2.split(cv2.cvtColor(image, cv2.COLOR_RGB2YCR_CB))[1]
        elif self.red_channel == 'U (YUV)':
            ret = cv2.split(cv2.cvtColor(image, cv2.COLOR_RGB2YUV))[1]
        return ret

    def _threshold_image(self, image):
        if self.threshold_enable:
            image = cv2.threshold(
                image, self.threshold_value, 255, cv2.THRESH_TOZERO)[1]
            if self.blur_enable:
                image = cv2.blur(image, (self.blur_value, self.blur_value))
            image = cv2.threshold(
                image, self.threshold_value, 255, cv2.THRESH_TOZERO)[1]
        return image

    def _window_mask(self, image):
        if self.window_enable:
            peak = image.argmax(axis=1)
            _min = peak - self.window_value
            _max = peak + self.window_value + 1
            mask = np.zeros_like(image)
            for i in xrange(self.calibration_data.height):
                mask[i, _min[i]:_max[i]] = 255
            # Apply mask
            image = cv2.bitwise_and(image, mask)
        return image

    # Segmented gaussian filter

    def _sgf(self, u, v, s):
        if len(u) > 1:
            i = 0
            sigma = 2.0
            f = np.array([])
            segments = [s[_r] for _r in np.ma.clump_unmasked(np.ma.masked_equal(s, 0))]
            # Detect stripe segments
            for segment in segments:
                j = len(segment)
                # Apply gaussian filter
                fseg = scipy.ndimage.gaussian_filter(u[i:i + j], sigma=sigma)
                f = np.concatenate((f, fseg))
                i += j
            return f, v
        else:
            return u, v

    # RANSAC implementation: https://github.com/ahojnnes/numpy-snippets/blob/master/ransac.py

    def _ransac(self, u, v):
        if len(u) > 1:
            data = np.vstack((v.ravel(), u.ravel())).T
            dr, thetar = self.ransac(data, self.LinearLeastSquares2D(), 2, 2)
            u = (dr - v * math.sin(thetar)) / math.cos(thetar)
        return u, v

    class LinearLeastSquares2D(object):
        '''
        2D linear least squares using the hesse normal form:
            d = x*sin(theta) + y*cos(theta)
        which allows you to have vertical lines.
        '''

        def fit(self, data):
            data_mean = data.mean(axis=0)
            x0, y0 = data_mean
            if data.shape[0] > 2:  # over determined
                u, v, w = np.linalg.svd(data - data_mean)
                vec = w[0]
                theta = math.atan2(vec[0], vec[1])
            elif data.shape[0] == 2:  # well determined
                theta = math.atan2(data[1, 0] - data[0, 0], data[1, 1] - data[0, 1])
            theta = (theta + math.pi * 5 / 2) % (2 * math.pi)
            d = x0 * math.sin(theta) + y0 * math.cos(theta)
            return d, theta

        def residuals(self, model, data):
            d, theta = model
            dfit = data[:, 0] * math.sin(theta) + data[:, 1] * math.cos(theta)
            return np.abs(d - dfit)

        def is_degenerate(self, sample):
            return False

    def ransac(self, data, model_class, min_samples, threshold, max_trials=100):
        '''
        Fits a model to data with the RANSAC algorithm.
        :param data: numpy.ndarray
            data set to which the model is fitted, must be of shape NxD where
            N is the number of data points and D the dimensionality of the data
        :param model_class: object
            object with the following methods implemented:
             * fit(data): return the computed model
             * residuals(model, data): return residuals for each data point
             * is_degenerate(sample): return boolean value if sample choice is
                degenerate
            see LinearLeastSquares2D class for a sample implementation
        :param min_samples: int
            the minimum number of data points to fit a model
        :param threshold: int or float
            maximum distance for a data point to count as an inlier
        :param max_trials: int, optional
            maximum number of iterations for random sample selection, default 100
        :returns: tuple
            best model returned by model_class.fit, best inlier indices
        '''

        best_model = None
        best_inlier_num = 0
        best_inliers = None
        data_idx = np.arange(data.shape[0])
        for _ in xrange(max_trials):
            sample = data[np.random.randint(0, data.shape[0], 2)]
            if model_class.is_degenerate(sample):
                continue
            sample_model = model_class.fit(sample)
            sample_model_residua = model_class.residuals(sample_model, data)
            sample_model_inliers = data_idx[sample_model_residua < threshold]
            inlier_num = sample_model_inliers.shape[0]
            if inlier_num > best_inlier_num:
                best_inlier_num = inlier_num
                best_inliers = sample_model_inliers
        if best_inliers is not None:
            best_model = model_class.fit(data[best_inliers])
        return best_model

# -*- coding: utf-8 -*-
# This file is part of the Horus Project


__author__ = 'Jes煤s Arroyo Torrens <jesus.arroyo@bq.com>'
__copyright__ = 'Copyright (C) 2014-2016 Mundo Reader S.L.'
__license__ = 'GNU General Public License v2 http://www.gnu.org/licenses/gpl2.html'


import cv2
import math
import numpy as np
import scipy.ndimage


from horus import Singleton
from horus.engine.calibration.calibration_data import CalibrationData##这是一个数据标定文件
from horus.engine.algorithms.point_cloud_roi import PointCloudROI##感兴趣区域的点云数据




@Singleton
class LaserSegmentation(object):


    def __init__(self):
        self.calibration_data = CalibrationData()
        self.point_cloud_roi = PointCloudROI()##对应整个工程中一个文件Point_cloud_roi.py，主要对激光照射的区域进行点云重建。


        self.red_channel = 'R (RGB)'##我的理解是激光束为红色，所以对红色通道开启。
        self.threshold_enable = False
        self.threshold_value = 0
        self.blur_enable = False
        self.blur_value = 0
        self.window_enable = False
        self.window_value = 0
        self.refinement_method = 'SGF'##seminiferous growth factor (SGF)增强现实框架，一种增强方式。


    def set_red_channel(self, value):
        self.red_channel = value


    def set_threshold_enable(self, value):
        self.threshold_enable = value


    def set_threshold_value(self, value):
        self.threshold_value = value


    def set_blur_enable(self, value):
        self.blur_enable = value


    def set_blur_value(self, value):
        self.blur_value = 2 * value + 1


    def set_window_enable(self, value):
        self.window_enable = value


    def set_window_value(self, value):
        self.window_value = value


    def set_refinement_method(self, value):
        self.refinement_method = value


    def compute_2d_points(self, image):##计算2D点云
        if image is not None:##如果图片非空
            image = self.compute_line_segmentation(image)
            # Peak detection: center of mass
            s = image.sum(axis=1)
            v = np.where(s > 0)[0]
            u = (self.calibration_data.weight_matrix * image).sum(axis=1)[v] / s[v]
            if self.refinement_method == 'SGF':
                # Segmented gaussian filter
                u, v = self._sgf(u, v, s)
            elif self.refinement_method == 'RANSAC':
                # Random sample consensus
                u, v = self._ransac(u, v)
            return (u, v), image


    def compute_hough_lines(self, image):
        if image is not None:
            image = self.compute_line_segmentation(image)
            lines = cv2.HoughLines(image, 1, np.pi / 180, 120)
            #函数HoughLines的原型为：
#void HoughLines(InputArray image,OutputArray lines, double rho, double theta, int threshold, double srn=0,double stn=0 )
#image为输入图像，要求是单通道的二值图像
#lines为输出直线向量，两个元素的向量（ρ，θ）代表一条直线，ρ是从原点（图像的左上角）的距离，θ是直线的角度（单位是弧度），0表示垂直线，π/2表示水平线
#rho为距离分辨率
#theta为角度分辨率
#threshold为阈值，在步骤2中，只有大于该值的点才有可能被当作极大值，即至少有多少条正弦曲线交于一点才被认为是直线
#srn和stn在多尺度霍夫变换的时候才会使用，在这里我们只研究经典霍夫变换
 
            # if lines is not None:
            #   rho, theta = lines[0][0]
            #   ## Calculate coordinates
            #   u1 = rho / np.cos(theta)
            #   u2 = u1 - height * np.tan(theta)
            return lines


    def compute_line_segmentation(self, image, roi_mask=False):
        if image is not None:
            # Apply ROI mask
            if roi_mask:
                image = self.point_cloud_roi.mask_image(image)##这是一个很重要的掩膜算法，就像是用一个
                #蒙皮与线色的激光线相与，可以把除激光线以外的图形挡住，只提取红色激光ROI区域。
            # Obtain red channel
            image = self._obtain_red_channel(image)##获得红色通道图形。见下面函数。。。。
            if image is not None:
                # Threshold image
                image = self._threshold_image(image)
                # Window mask
                image = self._window_mask(image)
            return image


    def _obtain_red_channel(self, image):
        ret = None
        if self.red_channel == 'R (RGB)':##通道分解 http://blog.csdn.net/eric_pycv/article/details/72887758
            ret = cv2.split(image)[0]
        elif self.red_channel == 'Cr (YCrCb)':##不同的图形空间的通道分解。
            ret = cv2.split(cv2.cvtColor(image, cv2.COLOR_RGB2YCR_CB))[1]
        elif self.red_channel == 'U (YUV)':##不同的图形空间的通道分解。
            ret = cv2.split(cv2.cvtColor(image, cv2.COLOR_RGB2YUV))[1]
        return ret


    def _threshold_image(self, image):
        if self.threshold_enable:
            image = cv2.threshold(
                image, self.threshold_value, 255, cv2.THRESH_TOZERO)[1]
            if self.blur_enable:
                image = cv2.blur(image, (self.blur_value, self.blur_value))
            image = cv2.threshold(
                image, self.threshold_value, 255, cv2.THRESH_TOZERO)[1]
        return image


        #cv2.threshold() 
#这个函数有四个参数，第一个原图像，第二个进行分类的阈值，第三个是高于（低于）阈值时赋予的新值，第四个是一个方法选择参数，常用的有： 
# cv2.THRESH_BINARY（黑白二值） 
# cv2.THRESH_BINARY_INV（黑白二值反转） 
# cv2.THRESH_TRUNC （得到的图像为多像素值） 
# cv2.THRESH_TOZERO 
# cv2.THRESH_TOZERO_INV 
#该函数有两个返回值，第一个retVal（得到的阈值值（在后面一个方法中会用到）），第二个就是阈值化后的图像。 
#http://blog.csdn.net/on2way/article/details/46812121
    def _window_mask(self, image):
        if self.window_enable:
            peak = image.argmax(axis=1)
            _min = peak - self.window_value
            _max = peak + self.window_value + 1
            mask = np.zeros_like(image)
            for i in xrange(self.calibration_data.height):
                mask[i, _min[i]:_max[i]] = 255
            # Apply mask
            image = cv2.bitwise_and(image, mask)
        return image


    # Segmented gaussian filter


    def _sgf(self, u, v, s):
        if len(u) > 1:
            i = 0
            sigma = 2.0
            f = np.array([])
            segments = [s[_r] for _r in np.ma.clump_unmasked(np.ma.masked_equal(s, 0))]
            # Detect stripe segments
            for segment in segments:
                j = len(segment)
                # Apply gaussian filter
                fseg = scipy.ndimage.gaussian_filter(u[i:i + j], sigma=sigma)
                f = np.concatenate((f, fseg))
                i += j
            return f, v
        else:
            return u, v


    # RANSAC implementation: https://github.com/ahojnnes/numpy-snippets/blob/master/ransac.py


    def _ransac(self, u, v):
        if len(u) > 1:
            data = np.vstack((v.ravel(), u.ravel())).T
            dr, thetar = self.ransac(data, self.LinearLeastSquares2D(), 2, 2)
            u = (dr - v * math.sin(thetar)) / math.cos(thetar)
        return u, v


    class LinearLeastSquares2D(object):
        '''
        2D linear least squares using the hesse normal form:
            d = x*sin(theta) + y*cos(theta)
        which allows you to have vertical lines.
        '''


        def fit(self, data):
            data_mean = data.mean(axis=0)
            x0, y0 = data_mean
            if data.shape[0] > 2:  # over determined
                u, v, w = np.linalg.svd(data - data_mean)
                vec = w[0]
                theta = math.atan2(vec[0], vec[1])
            elif data.shape[0] == 2:  # well determined
                theta = math.atan2(data[1, 0] - data[0, 0], data[1, 1] - data[0, 1])
            theta = (theta + math.pi * 5 / 2) % (2 * math.pi)
            d = x0 * math.sin(theta) + y0 * math.cos(theta)
            return d, theta


        def residuals(self, model, data):
            d, theta = model
            dfit = data[:, 0] * math.sin(theta) + data[:, 1] * math.cos(theta)
            return np.abs(d - dfit)


        def is_degenerate(self, sample):
            return False
# RANSAC是“RANdom SAmple Consensus（随机抽样一致）”的缩写。
#它可以从一组包含“局外点”的观测数据集中，通过迭代方式估计数学模型的参数。
#它是一种不确定的算法——它有一定的概率得出一个合理的结果；为了提高概率必须提高迭代次数。
# 该算法最早由Fischler和Bolles于1981年提出?
#我的理解是，利用激光线会生成很多线状点云，但是也会有很多离散的点，那么利用RANSAC这个算法
#把线状点云集中在一条直线上。滤除掉那些离散的点。请参考如下链接。
#https://www.cnblogs.com/xrwang/archive/2011/03/09/ransac-1.html
    def ransac(self, data, model_class, min_samples, threshold, max_trials=100):
        '''
        Fits a model to data with the RANSAC algorithm.
        :param data: numpy.ndarray
            data set to which the model is fitted, must be of shape NxD where
            N is the number of data points and D the dimensionality of the data
        :param model_class: object
            object with the following methods implemented:
             * fit(data): return the computed model
             * residuals(model, data): return residuals for each data point
             * is_degenerate(sample): return boolean value if sample choice is
                degenerate
            see LinearLeastSquares2D class for a sample implementation
        :param min_samples: int
            the minimum number of data points to fit a model
        :param threshold: int or float
            maximum distance for a data point to count as an inlier
        :param max_trials: int, optional
            maximum number of iterations for random sample selection, default 100
        :returns: tuple
            best model returned by model_class.fit, best inlier indices
        '''


        best_model = None
        best_inlier_num = 0
        best_inliers = None
        data_idx = np.arange(data.shape[0])
        for _ in xrange(max_trials):
            sample = data[np.random.randint(0, data.shape[0], 2)]
            if model_class.is_degenerate(sample):
                continue
            sample_model = model_class.fit(sample)
            sample_model_residua = model_class.residuals(sample_model, data)
            sample_model_inliers = data_idx[sample_model_residua < threshold]
            inlier_num = sample_model_inliers.shape[0]
            if inlier_num > best_inlier_num:
                best_inlier_num = inlier_num
                best_inliers = sample_model_inliers
        if best_inliers is not None:
            best_model = model_class.fit(data[best_inliers])
        return best_model