Perception

def __init__(self, agent: Agent):
    self.agent = agent
    self.logger = logging.getLogger("Base Detector")

@abstractmethod
def run_step(self) -> Any:
    return None

def __init__(self, sky_level=0.9, t=0.05, del_ang=0.2, fit_type='exp', **kwargs):
    super().__init__(**kwargs)
    self.logger = logging.getLogger("Ground Plane Detector")
    self.sky_level = sky_level
    self.thresh = t
    self.fit_type = fit_type
    self.del_ang = del_ang
    self.roll_ang = 0
    self.rot_axis = [0, 0, 1]

    self.orig_preds = None
    self.preds = None

    self.curr_segmentation: Optional[np.ndarray] = None

    self.logger.info("Ground Plane Detector Initiated")

@staticmethod
def convert_to_log(x):
    return np.clip(1 + np.log(x + 1e-10) / 5.70378, 0.005, 1.0)

def get_roll_stats(self, depth_image):
    xyz = self.img_to_world(depth_image).T
    xyz_samp = xyz[np.random.choice(xyz.shape[0], 400, replace=False), :]
    u, s, vt = np.linalg.svd(xyz_samp - xyz_samp.mean())

    reg = vt[2]
    no_y = np.array(reg, copy=True)
    no_y[1] = 0
    nvt, nx = np.linalg.norm(reg), np.linalg.norm(no_y)
    cos_ang = np.dot(reg, no_y) / (nvt * nx)

    unitcross = lambda a, b: np.cross(a, b) / np.linalg.norm(np.cross(a, b))
    rot_axis = unitcross(vt[2], no_y)

    return np.arccos(cos_ang), rot_axis  # radians

def gpd_mesh(self, depth_image):
    xs = []
    data = []
    max_depth = 0
    for i in range(depth_image.shape[0] - 1, -1, -1):
        j = np.argmax(depth_image[i, :])
        d = depth_image[i][j]
        if d > 0.3:
            break
        if d > max_depth and d > 0.01:
            max_depth = d
            xs.append(i)
            data.append(d)

    xs = np.array(xs[::-1], dtype=np.float64)
    data = np.array(data[::-1], dtype=np.float64)

    if self.fit_type == 'lsq':
        a, b, c, d = _Leastsq_Exp.fit(xs / xs.max(), data)
        pred_func = _Leastsq_Exp.construct_f(a, b, c, d)
        rows = np.meshgrid(
            np.arange(depth_image.shape[1]), np.arange(depth_image.shape[0])
        )[1]
        preds = pred_func(rows / rows.max())
        preds[preds > 1] = 0
        return preds
    else:
        a, b, c, p, q = _Exponential_Model.fit(xs, data)
        pred_func = _Exponential_Model.construct_f(a, b, c, p, q)
        rows = np.meshgrid(
            np.arange(depth_image.shape[1]), np.arange(depth_image.shape[0])
        )[1]
        preds = pred_func(rows)
        preds[preds > 1] = 0
        return preds

def img_to_world(self, depth_img, sky_level=0.2, depth_scaling_factor=1000) -> np.ndarray:
    # (Intrinsic) K Matrix
    intrinsics_matrix = self.agent.front_depth_camera.intrinsics_matrix

    # get a 2 x N array for their indices
    bool_mat3 = (0.1 < depth_img) * (depth_img < sky_level)
    # bool_mat2 = 0.1 < depth_img
    # bool_mat3 = bool_mat * bool_mat2
    ground_loc = np.where(bool_mat3)
    depth_val = depth_img[bool_mat3] * depth_scaling_factor
    ground_loc = ground_loc * depth_val

    # compute raw_points
    raw_points = np.vstack([ground_loc, depth_val])

    # convert to cords_y_minus_z_x
    return np.linalg.inv(intrinsics_matrix) @ raw_points

def output_gpd(self, d_frame):
    # first im going to find out where is the sky
    sky = np.where(d_frame > self.sky_level)

    # then im going to find out where is the ground
    ground = np.where(np.abs(d_frame - self.preds) < self.thresh)

    result = np.zeros(shape=(d_frame.shape[0], d_frame.shape[1], 3))
    result[ground] = self.GROUND
    result[sky] = self.SKY
    # result = result.astype('uint8')
    # # try:
    # #     new_roll_ang, self.rot_axis = self.get_roll_stats(d_frame)  # this method is pretty slow
    # #     if np.abs(self.roll_ang - new_roll_ang) > self.del_ang:
    # #         print(f"Recalibrating {self.agent.time_counter}")
    # #         self.roll_ang = new_roll_ang
    # #         self.preds = self.roll_frame(self.orig_preds, self.roll_ang, -1 * self.rot_axis)
    # # except Exception as e:
    # #     self.logger.error(f"Failed to compute output: {e}")
    #
    # result = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
    # # img = cv2.threshold(result, 127, 255, cv2.THRESH_BINARY)[1]
    # ret, thresh = cv2.threshold(result, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # # print(np.shape(img))
    # # num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(image=thresh, connectivity=8,ltype=cv2.CV_16U) # img= (600 x 800)
    # _, labels = cv2.connectedComponents(image=thresh)
    # indicies = np.where(labels == 2)
    # print(np.shape(indicies))
    # #
    cv2.imshow("segmented", result)
    cv2.waitKey(1)
    return result

def reg_img_to_world(self, depth_img, sky_level=0.9, depth_scaling_factor=1000) -> np.ndarray:
    # (Intrinsic) K Matrix

    intrinsics_matrix = self.agent.front_depth_camera.intrinsics_matrix

    # get a 2 x N array for their indices
    bool_mat = depth_img > (depth_img.min() - 1)
    ground_loc = np.where(bool_mat)
    depth_val = depth_img[bool_mat] * depth_scaling_factor
    ground_loc = ground_loc * depth_val

    # compute raw_points
    raw_points = np.vstack([ground_loc, depth_val])

    return np.linalg.inv(intrinsics_matrix) @ raw_points

def roll_frame(self, depth_image, ang, rot_axis, no_axis=False):
    if no_axis:
        return depth_image

    xyz = self.reg_img_to_world(depth_image).T
    xyz_mean = xyz.mean(axis=0)
    xyz = xyz - xyz_mean

    cos_ang = np.cos(ang)
    x, y, z = rot_axis

    c = cos_ang
    s = np.sqrt(1 - c * c)
    C = 1 - c
    rmat = np.array([[x * x * C + c, x * y * C - z * s, x * z * C + y * s],
                     [y * x * C + z * s, y * y * C + c, y * z * C - x * s],
                     [z * x * C - y * s, z * y * C + x * s, z * z * C + c]])
    rot_xyz = rmat @ (xyz.T)
    rot_xyz = rot_xyz.T + xyz_mean
    return rot_xyz[:, 2].reshape(depth_image.shape) / 1000

def run_step(self) -> Any:
    logged_depth = self.convert_to_log(self.agent.front_depth_camera.data.copy())
    if self.orig_preds is None or self.preds is None:
        self.orig_preds = self.gpd_mesh(logged_depth)
        self.preds = np.copy(self.orig_preds)
        self.logger.debug("Ground Plane Preds Computed")
    else:
        self.curr_segmentation = self.output_gpd(logged_depth)

def __init__(self,
             vehicle: Vehicle,
             camera: Camera,
             sky_line_level: int = 310,
             max_detectable_distance_threshold: float = 0.08,
             min_caliberation_boundary: float = 0.01):
    super().__init__(vehicle=vehicle, camera=camera)
    self.logger = logging.getLogger(__name__)
    self._sky_line_level = sky_line_level
    self._max_detectable_distance_threshold = max_detectable_distance_threshold
    self._min_caliberation_boundary = min_caliberation_boundary
    self._test_depth_img: Optional[np.array] = None
    self._predict_matrix: Optional[np.array] = None  # preds

    self.semantic_segmentation: Optional[np.array] = None
    self.curr_ground: Optional[np.array] = None

    self.logger.debug("Ground Plane Detector Initialized")

def construct_f(self, a, b, c, p, q):
    def f(x):
        # print(type(a), type(b), type(c), type(p), type(q))
        return a + b * np.exp(p * x) + c * np.exp(q * x)

    return f

def F1(self, SS_k, y_k):
    return SS_k / y_k

def F2(self, S_k, y_k):
    return S_k / y_k

def F3(self, x_k, y_k):
    return (x_k ** 2) / y_k

def F4(self, x_k, y_k):
    return x_k / y_k

def F5(self, y_k):
    return 1 / y_k

def fit(self, x, y):
    Sxy = self.S(x, y)
    SSxy = self.SS(Sxy, x)
    F1xy = self.F1(SSxy, y)
    F2xy = self.F2(Sxy, y)
    F3xy = self.F3(x, y)
    F4xy = self.F4(x, y)
    F5xy = self.F5(y)
    F = np.array([F1xy, F2xy, F3xy, F4xy, F5xy])
    f = np.array([np.sum(F1xy), np.sum(F2xy), np.sum(F3xy),
                  np.sum(F4xy), np.sum(F5xy)])
    F = F @ F.T
    A, B, C, D, E = np.linalg.inv(F) @ f
    pre_sqrt = np.clip(B ** 2 + 4 * A, 0, np.inf)  # edits 1

    p = 0.5 * (B + np.sqrt(pre_sqrt))
    q = 0.5 * (B - np.sqrt(pre_sqrt))
    G1 = 1 / y
    G2 = np.exp(p * x) / y
    G3 = np.exp(q * x) / y
    G = np.array([G1, G2, G3])
    G = G @ G.T
    g = np.array([np.sum(G1), np.sum(G2), np.sum(G3)])
    a, b, c = np.linalg.pinv(G) @ g  # edits 2
    return a, b, c, p, q

def recalibrate(self,
                sky_line_level: int = 310,
                max_detectable_distance_threshold: float = 0.089,
                min_caliberation_boundary: float = 0.01):
    """
    Force a recalibration of the ground plane prediction matrix
    Returns:
        None
    """
    self._test_depth_img = None
    self._predict_matrix = None

    self._sky_line_level = sky_line_level
    self._max_detectable_distance_threshold = max_detectable_distance_threshold
    self._min_caliberation_boundary = min_caliberation_boundary

def run_step(self, vehicle: Vehicle, new_data: np.array):
    """
    This function assumes that the function calling it will set the variable self.curr_depth_img
    In the first run of this function,
        it will remember the input depth image as self._test_depth_img
    In the second run of this function,
        it will calculate the prediction matrix
    In the preceding runs, it will use the prediction matrix to find ground plane


    Args:
        vehicle: current vehicle state
        new_data: current frame for this detector

    Returns:
        None

    """
    super(SemanticSegmentationDetector, self).run_step(vehicle, new_data)
    if self._test_depth_img is None:
        self._test_depth_img = png_to_depth(new_data)
        return
    elif self._predict_matrix is None:
        # try calibrate on the second frame received
        xs = []
        data = []
        depth_array = png_to_depth(new_data)
        # depth_image = calibration image, grab from somewhere

        for i in range(self._sky_line_level+10, depth_array.shape[0]):
            j = np.argmax(depth_array[i, :])

            if depth_array[i][j] > self._min_caliberation_boundary:
                xs.append(i)
                data.append(depth_array[i][j])
        a, b, c, p, q = self.fit(
            np.array(xs, dtype=np.float64),
            np.array(data, dtype=np.float64)
        )
        test_image = self._test_depth_img
        pred_func = self.construct_f(a, b, c, p, q)
        rows = np.meshgrid(
            np.arange(test_image.shape[1]), np.arange(test_image.shape[0])
        )[1]
        self._predict_matrix = pred_func(rows)
        return
    else:
        depth_array = png_to_depth(new_data.copy())  # this turns it into 2D np array of shape (Width x Height)
        semantic_seg = np.zeros(shape=np.shape(new_data))

        # find sky and ground
        sky = np.where(depth_array == 1)
        ground = np.where(np.abs(depth_array - self._predict_matrix) > self._max_detectable_distance_threshold)

        semantic_seg[ground] = [255, 255, 255]
        semantic_seg[sky] = [255, 0, 0]  # BGR???
        self.semantic_segmentation = semantic_seg

def S(self, x, y):
    ret_S = [0]
    for k in range(1, len(x)):
        S_k = self.Sk(ret_S[k - 1], y[k], y[k - 1], x[k], x[k - 1])
        ret_S.append(S_k)
    return ret_S

def Sk(self, S_k_1, y_k, y_k_1, x_k, x_k_1):
    return S_k_1 + 0.5 * (y_k + y_k_1) * (x_k - x_k_1)

def SS(self, s, x):
    ret_SS = [0]
    for k in range(1, len(x)):
        SS_k = self.SSk(ret_SS[k - 1], s[k], s[k - 1], x[k], x[k - 1])
        ret_SS.append(SS_k)
    return ret_SS

def SSk(self, SS_k_1, S_k, S_k_1, x_k, x_k_1):
    return SS_k_1 + 0.5 * (S_k + S_k_1) * (x_k - x_k_1)

def __init__(self, max_detectable_distance=0.1, depth_scaling_factor=1000, **kwargs):
    """

    Args:
        max_detectable_distance: maximum detectable distance in km
        depth_scaling_factor: scaling depth back to world scale. 1000 m = 1 km
        **kwargs:
    """
    super().__init__(**kwargs)
    self.max_detectable_distance = max_detectable_distance
    self.depth_scaling_factor = depth_scaling_factor
    self.logger = logging.getLogger("Point Cloud Detector")
    self.pcd: o3d.geometry.PointCloud = o3d.geometry.PointCloud()
    self.vis = o3d.visualization.Visualizer()

    self.counter = 0

def calculate_world_cords(self, max_points_to_convert=5000):
    depth_img = self.agent.front_depth_camera.data.copy()
    depth_img[depth_img > 0.05] = 0
    depth_img = depth_img * 1000
    # create a n x 2 array
    img_pos = np.indices((depth_img.shape[0], depth_img.shape[1])).transpose(1,2,0)
    img_pos = np.reshape(a=img_pos, newshape=(img_pos.shape[0] * img_pos.shape[1], 2))
    depth_array = np.reshape(a=depth_img, newshape=(depth_img.shape[0] * depth_img.shape[1], 1))

    # print(self.agent.front_depth_camera.data[0][0], self.agent.front_depth_camera.data[0][2])

    """
    [[x,y, depth]]
    """
    # indicies = np.random.choice(a=len(img_pos), size=max_points_to_convert, replace=False)

    # selected_depth_array = depth_array[indicies, :]
    # selected_img_pos = img_pos[indicies, :] * selected_depth_array
    selected_depth_array = depth_array
    selected_img_pos = img_pos * selected_depth_array
    #
    raw_p2d = np.append(selected_img_pos, selected_depth_array, axis=1)

    #
    # raw_p2d = raw_p2d[raw_p2d[:, 2] > 0]
    # raw_p2d = raw_p2d[raw_p2d[:, 2] < 20]
    #
    # # convert to cords_y_minus_z_x
    cords_y_minus_z_x = np.linalg.inv(self.agent.front_depth_camera.intrinsics_matrix) @ raw_p2d.T
    # #
    # # convert to cords_xyz_1 -y,x-z
    cords_xyz_1 = np.vstack([
        -cords_y_minus_z_x[0, :],
        cords_y_minus_z_x[2, :],
        -cords_y_minus_z_x[1, :],
        np.ones((1, np.shape(cords_y_minus_z_x)[1]))
    ])
    print(self.agent.vehicle.transform)
    self.agent.vehicle.transform.rotation.roll = 90  #-self.agent.vehicle.transform.rotation.roll
    # self.agent.vehicle.transform.rotation.pitch = 90
    self.agent.vehicle.transform.rotation.yaw = -self.agent.vehicle.transform.rotation.yaw
    # # print(self.agent.front_depth_camera.transform)
    # # self.agent.front_depth_camera.transform.rotation.roll = -self.agent.front_depth_camera.transform.rotation.roll
    # # self.agent.front_depth_camera.transform.rotation.pitch = -self.agent.front_depth_camera.transform.rotation.pitch
    # self.agent.front_depth_camera.transform.rotation.yaw = -self.agent.front_depth_camera.transform.rotation.yaw
    #
    points = self.agent.vehicle.transform.get_matrix() @ \
             self.agent.front_depth_camera.transform.get_matrix() @ cords_xyz_1
    return points[:3, :]

def run_step(self) -> Optional[np.ndarray]:
    points_3d = self.calculate_world_cords(max_points_to_convert=10000)

    # print(np.amin(points_3d, axis=1), np.amax(points_3d, axis=1), self.agent.vehicle.transform.location)
    # self.pcd.points = o3d.utility.Vector3dVector(points_3d.T - np.mean(points_3d.T, axis=0))
    # self.pcd.paint_uniform_color([0,0,0])
    # points_3d = points_3d.T
    # # filter out anything that is "above" my vehicle (so ground is definitely below my vehicle)
    # # this part is shady, idk why it works
    # ground_indicies = np.where(points_3d[:, 0] > self.agent.vehicle.transform.location.to_array()[0])
    # ground_points = points_3d[ground_indicies]
    # print(np.shape(ground_points.T))
    # self.pcd.points = o3d.utility.Vector3dVector(ground_points.T)
    #
    # # # find the normals using Open3D
    # self.pcd.points = o3d.utility.Vector3dVector(ground_points - np.mean(ground_points, axis=0))
    # print(self.pcd.get_min_bound(), self.pcd.get_max_bound())
    # self.pcd.estimate_normals(fast_normal_computation=True)
    #
    # # find normals that are less than the mean normal,
    # # since I know that most of the things in front of me are going to be ground
    # normals = np.asarray(self.pcd.normals)
    # abs_diff = np.linalg.norm(normals - np.mean(normals, axis=0), axis=1)
    # ground_loc = np.where(abs_diff < np.mean(abs_diff))
    # ground = ground_points[ground_loc[0]]
    #
    # # turn it into Open3D PointCloud object again to utilize its remove outlier method
    # # this is when I drive to the side, the opposing road will be recognized, but we don't want that
    # self.pcd.points = o3d.utility.Vector3dVector(ground)
    # new_pcd, indices = self.pcd.remove_radius_outlier(100, 2)

    # project it back to Open3D PointCloud object for visualizations
    # minus the mean for stationary visualization
    # # self.pcd.points = o3d.utility.Vector3dVector(ground[indices])
    # if self.counter == 0:
    #     self.vis.create_window(window_name="Open3d", width=400, height=400)
    #     self.vis.add_geometry(self.pcd)
    #     render_option: o3d.visualization.RenderOption = self.vis.get_render_option()
    #     render_option.show_coordinate_frame = True
    # else:
    #     self.vis.update_geometry(self.pcd)
    #     render_option: o3d.visualization.RenderOption = self.vis.get_render_option()
    #     render_option.show_coordinate_frame = True
    #     self.vis.poll_events()
    #     self.vis.update_renderer()
    # self.counter += 1
    return points_3d