diff --git a/persontracking/algorithms.py b/persontracking/algorithms.py new file mode 100644 index 0000000..a968b14 --- /dev/null +++ b/persontracking/algorithms.py @@ -0,0 +1,107 @@ +import math + +class Group: + + def __init__(self, number, points=[]): + self._points = points + + def add_point(self, point): + self._points.append(point) + self._update_min_max(point) + + def get_points(self): + return self._points + + # Should use property here. + def set_number(self, number): + self._number = number + + def get_number(self): + return self._number + + def _update_min_max(self, new_point): + ''' + Updates the in and max points for this group. + This is to determine when assigning groups whether the + same group is selected. + ''' + converted_point = self.convert_lidar_to_cartesian(new_point) + + if not hasattr(self, '_minX') or self._minX > converted_point[0]: + self._minX = converted_point[0] + + if not hasattr(self, '_maxX') or self._maxX < converted_point[0]: + self._maxX = converted_point[0] + + if not hasattr(self, '_miny') or self._miny > converted_point[1]: + self._minY = converted_point[1] + + if not hasattr(self, '_maxY') or self._maxY < converted_point[1]: + self._maxY = converted_point[0] + + + def convert_lidar_to_cartesian(self, new_point): + x = new_point[2] * math.sin(new_point[1]) + y = new_point[2] * math.cos(new_point[1]) + return (x,y) + + def get_minX(self): + return self._minY + + def get_maxX(self): + return self._maxY + + def get_minY(self): + return self._minY + + def get_maxY(self): + return self._maxY + +def calc_groups(scan): + ''' + Calculates groups of points from a lidar scan. The scan should + already be sorted. + Should return all groups. + ''' + prevPoint = None + currentGroup = None + allGroups = [] + currentGroupNumber = 0 + + # assume the list is already sorted. + for point in scan: + if prevPoint is None: + prevPoint = point + continue + + # Distances are in mm. + if (point[2] - prevPoint[2]) ** 2 < 10 ** 2: # within 1cm makes a group. Will need to play around with this. + if currentGroup is None: + currentGroup = Group(currentGroupNumber) + allGroups.append(currentGroup) + currentGroup.add_point(point) + else: + if currentGroup is not None: + currentGroupNumber += 1 + currentGroup = None + + prevPoint = point + + return allGroups + +def find_centre(group): + return ((group.get_maxX() + group.get_minX()) / 2, (group.get_maxY() + group.get_minY()) / 2) + +def assign_groups(prev_groups, new_groups): + ''' + Assigns group numbers to a new scan based on the groups of an old scan. + ''' + for group in prev_groups: + old_centre = find_centre(prev_groups) + for new_group in new_groups: + new_centre = find_centre(new_group) + # They are considered the same if the new group and old group centres are within 5cm. + if ((new_centre[0] - old_centre[0]) ** 2 + (new_centre[1] - old_centre[1]) ** 2) < 50 ** 2: + new_group.set_number(group.get_number()) + + return new_groups \ No newline at end of file