refactor

skais/ais/ais_trajectory.py

@@ -61,11 +61,54 @@ def apply_time_sequence(dat, time, func):
 
 
 def __get_image_value__(features, bounds):
+    if len(bounds) < 1:
+        return [1]
     value = []
     for f, b in zip(features, bounds):
         value.append(1 - (b[1] - f - b[0]) / (b[1] - b[0]))
     return value
 
+def __get_bounding_box__(bounding_box, positions, ref_index):
+    if bounding_box == 'fit':
+        lower_lon, upper_lon = (min(positions[:, 0]), max(positions[:, 0]))
+        lower_lat, upper_lat = (min(positions[:, 1]), max(positions[:, 1]))
+    elif bounding_box == 'centered':
+        center_lon, center_lat = positions[ref_index]
+        min_lon, max_lon = (min(positions[:, 0]), max(positions[:, 0]))
+        min_lat, max_lat = (min(positions[:, 1]), max(positions[:, 1]))
+
+        distance_to_center = max(center_lon - min_lon, max_lon - center_lon, center_lat - min_lat,
+                                 max_lat - center_lat)
+
+        upper_lat = center_lat + distance_to_center
+        lower_lat = center_lat - distance_to_center
+        upper_lon = center_lon + distance_to_center
+        lower_lon = center_lon - distance_to_center
+    elif type(bounding_box) is list:
+        if type(bounding_box[0]) is not numbers.Number:
+            upper_lon = bounding_box[1][0]
+            lower_lon = bounding_box[0][0]
+            upper_lat = bounding_box[1][1]
+            lower_lat = bounding_box[0][1]
+        else:
+            center_lon, center_lat = positions[ref_index]
+            distance_to_center_lon = bounding_box[0]
+            distance_to_center_lat = bounding_box[1]
+            upper_lat = center_lat + distance_to_center_lat
+            lower_lat = center_lat - distance_to_center_lat
+            upper_lon = center_lon + distance_to_center_lon
+            lower_lon = center_lon - distance_to_center_lon
+    else:
+        raise ValueError(f"Option not supported: {bounding_box}")
+
+    if lower_lat == upper_lat:
+        lower_lat -= 1
+        upper_lat += 1
+    if lower_lon == upper_lon:
+        lower_lon -= 1
+        upper_lon += 1
+
+    return lower_lon, upper_lon, lower_lat, upper_lat
 
 class AISTrajectory(AISPoints):
     def __init__(self, df, mmsi=0, interpolation_time=None):
@@ -233,52 +276,40 @@ class AISTrajectory(AISPoints):
                 result.append((row['ts_sec'], current_label))
         return result
 
-    def generate_array_from_positions(self, height=256, width=256, link=True, bounding_box='fit', ref_index=-1, features=None,
-                                      node_size=0):
-        nb_channels = 1
+
+
+    def generate_array_from_positions(self, height=256, width=256, link=True, bounding_box='fit', ref_index=-1,
+                                      features=None, node_size=0):
 
         positions = self.df[['longitude', 'latitude']].to_numpy()
-        if bounding_box == 'fit':
-            lower_lon, upper_lon = (min(positions[:, 0]), max(positions[:, 0]))
-            lower_lat, upper_lat = (min(positions[:, 1]), max(positions[:, 1]))
-        elif bounding_box == 'centered':
-            center_lon, center_lat = positions[ref_index]
-            min_lon, max_lon = (min(positions[:, 0]), max(positions[:, 0]))
-            min_lat, max_lat = (min(positions[:, 1]), max(positions[:, 1]))
 
-            distance_to_center = max(center_lon - min_lon, max_lon - center_lon, center_lat - min_lat,
-                                     max_lat - center_lat)
+        lower_lon, upper_lon, lower_lat, upper_lat = __get_bounding_box__(bounding_box, positions, ref_index)
 
-            upper_lat = center_lat + distance_to_center
-            lower_lat = center_lat - distance_to_center
-            upper_lon = center_lon + distance_to_center
-            lower_lon = center_lon - distance_to_center
-        elif type(bounding_box) is list:
-            if type(bounding_box[0]) is not numbers.Number:
-                upper_lon = bounding_box[1][0]
-                lower_lon = bounding_box[0][0]
-                upper_lat = bounding_box[1][1]
-                lower_lat = bounding_box[0][1]
+        bounds = []
+        if features is None:
+            features_vectors = None
+        elif type(features) is list:
+            features_vectors = self.df[features].to_numpy()
+            for c in features_vectors.T:
+                bounds.append((0, max(c)))
+        elif type(features) is str:
+            features_vectors = self.df[[features]].to_numpy()
+            for c in features_vectors.T:
+                bounds.append((0, max(c)))
+        elif type(features) is dict:
+            bounds = list(features.values())
+            features_vectors = self.df[features.keys()].to_numpy()
         else:
-                center_lon, center_lat = positions[ref_index]
-                distance_to_center_lon = bounding_box[0]
-                distance_to_center_lat = bounding_box[1]
-                upper_lat = center_lat + distance_to_center_lat
-                lower_lat = center_lat - distance_to_center_lat
-                upper_lon = center_lon + distance_to_center_lon
-                lower_lon = center_lon - distance_to_center_lon
+            raise TypeError("Type not supported")
+
+        if features_vectors is not None:
+            nb_channels = len(features_vectors.T)
         else:
-            raise ValueError(f"Option not supported: {bounding_box}")
+            nb_channels = 1
+        data = np.zeros((height, width, nb_channels), dtype=np.float)
 
-        if lower_lat == upper_lat:
-            lower_lat -= 1
-            upper_lat += 1
-        if lower_lon == upper_lon:
-            lower_lon -= 1
-            upper_lon += 1
+        if features_vectors is None:
 
-        if features is None:
-            data = np.zeros((height, width, nb_channels), dtype=np.uint8)
             for longitude, latitude in positions:
                 x_coord, y_coord = get_coord(latitude, longitude, height, width, lower_lat, upper_lat, lower_lon,
                                              upper_lon)
@@ -302,27 +333,7 @@ class AISTrajectory(AISPoints):
                     lon, lat = longitude, latitude
                     for x, y in bresenham(x_prv, y_prev, x_nxt, y_nxt):
                         data[x, y] = [1]
-
-        else:
-            bounds = []
-            if type(features) is list:
-                features_vectors = self.df[features].to_numpy()
-                for c in features_vectors.T:
-                    bounds.append((0, max(c)))
-            elif type(features) is str:
-                features_vectors = self.df[[features]].to_numpy()
-                for c in features_vectors.T:
-                    bounds.append((0, max(c)))
-            elif type(features) is dict:
-                bounds = list(features.values())
-                features_vectors = self.df[features.keys()].to_numpy()
         else:
-                raise TypeError("Type not supported")
-
-            nb_channels = len(features_vectors.T)
-            data = np.zeros((height, width, nb_channels), dtype=np.float)
-
-
             for pos, f in zip(positions, features_vectors):
                 latitude = pos[1]
                 longitude = pos[0]