feat: build speedzone models

tf_container/train.py

@@ -15,7 +15,6 @@ import zipfile
 
 import tensorflow.keras.losses
 
-# from tensorflow.keras import backend as K
 from tensorflow.keras import callbacks
 from tensorflow.keras.layers import Conv2D
 from tensorflow.keras.layers import Dropout, Flatten, Dense
@@ -29,6 +28,15 @@ MODEL_CATEGORICAL = "categorical"
 MODEL_LINEAR = "linear"
 
 
+def linear_bin_speed_zone(a: int, N: int = 4) -> npt.NDArray[np.float64]:
+    """
+    create a bin of length N
+    """
+    arr = np.zeros(N)
+    arr[a] = 1
+    return arr
+
+
 def linear_bin(a: float, N: int = 15, offset: int = 1, R: float = 2.0) -> npt.NDArray[np.float64]:
     """
     create a bin of length N
@@ -57,6 +65,12 @@ def get_data(root_dir: pathlib.Path, filename: str) -> typing.List[typing.Any]:
     return [(d['user/angle']), root_dir, d['cam/image_array']]
 
 
+def get_data_speed_zone(root_dir, filename):
+    print('load data from file ' + filename)
+    d = json.load(open(os.path.join(root_dir, filename)))
+    return [(d['speed_zone']), root_dir, d['cam/image_array']]
+
+
 numbers = re.compile(r'(\d+)')
 
 
@@ -66,8 +80,8 @@ def unzip_file(root: pathlib.Path, f: str) -> None:
     zip_ref.close()
 
 
-def train(model_type: str, batch_size: int, slide_size: int, img_height: int, img_width: int, img_depth: int,
+def train(model_type: str, record_field: str, batch_size: int, slide_size: int, img_height: int, img_width: int,
-          horizon: int, drop: float) -> None:
+          img_depth: int, horizon: int, drop: float) -> None:
     # env = cs.TrainingEnvironment()
 
     os.system('mkdir -p logs')
@@ -83,15 +97,24 @@ def train(model_type: str, batch_size: int, slide_size: int, img_height: int, im
             if f.endswith('.zip'):
                 unzip_file(pathlib.Path(root), f)
 
-    for root, dirs, files in os.walk('/opt/ml/input/data/train'):
+    if record_field == 'angle':
-        data.extend(
+        output_name = 'angle_out'
-            [get_data(pathlib.Path(root), f) for f in sorted(files, key=str.lower) if
+        for root, dirs, files in os.walk('/opt/ml/input/data/train'):
-             f.startswith('record') and f.endswith('.json')])
+            data.extend(
+                [get_data(root, f) for f in sorted(files, key=str.lower) if f.startswith('record') and f.endswith('.json')])
+    elif record_field == 'speed_zone':
+        output_name = 'speed_zone_output'
+        for root, dirs, files in os.walk('/opt/ml/input/data/train'):
+            data.extend(
+                [get_data_speed_zone(root, f) for f in sorted(files, key=str.lower) if f.startswith('record') and f.endswith('.json')])
+    else:
+        print(f"invalid record filed: {record_field}")
+        return
 
-    # ### Loading throttle and angle ###
+    # ### Loading values (angle or speed_zone) ###
 
-    angle = [d[0] for d in data]
+    value = [d[0] for d in data]
-    angle_array = np.array(angle)
+    value_array = np.array(value)
 
     # ### Loading images ###
     if horizon > 0:
@@ -104,7 +127,7 @@ def train(model_type: str, batch_size: int, slide_size: int, img_height: int, im
     # slide images vs orders
     if slide_size > 0:
         images = images[:len(images) - slide_size]
-        angle_array = angle_array[slide_size:]
+        value_array = value_array[slide_size:]
 
     # ### Start training ###
     from datetime import datetime
@@ -123,19 +146,28 @@ def train(model_type: str, batch_size: int, slide_size: int, img_height: int, im
     model_filepath = '/opt/ml/model/model_other'
     if model_type == MODEL_CATEGORICAL:
         model_filepath = '/opt/ml/model/model_cat'
-        angle_cat_array = np.array([linear_bin(float(a)) for a in angle_array])
+        if record_field == 'angle':
-        model = default_categorical(input_shape=(img_height - horizon, img_width, img_depth), drop=drop)
+            input_value_array = np.array([linear_bin(float(a)) for a in value_array])
-        loss = {'angle_out': 'categorical_crossentropy', }
+            output_bin = 15
+        elif record_field == 'speed_zone':
+            input_value_array = np.array([linear_bin_speed_zone(a) for a in value_array])
+            output_bin = 4
+        model = default_categorical(input_shape=(img_height - horizon, img_width, img_depth), drop=drop,
+                                    output_name=output_name, output_bin=output_bin)
+        loss = {output_name: 'categorical_crossentropy', }
         optimizer = 'adam'
     elif model_type == MODEL_LINEAR:
         model_filepath = '/opt/ml/model/model_lin'
-        angle_cat_array = np.array([a for a in angle_array])
+        input_value_array = np.array([a for a in value_array])
-        model = default_linear(input_shape=(img_height - horizon, img_width, img_depth), drop=drop)
+        model = default_linear(input_shape=(img_height - horizon, img_width, img_depth), drop=drop, output_name=output_name)
-        loss = tensorflow.keras.losses.Loss('mse')
+        loss = 'mse'
         optimizer = 'rmsprop'
     else:
         raise Exception("invalid model type")
 
+    # Display the model's architecture
+    model.summary()
+
     save_best = callbacks.ModelCheckpoint(model_filepath, monitor='val_loss', verbose=1,
                                           save_best_only=True, mode='min')
     early_stop = callbacks.EarlyStopping(monitor='val_loss',
@@ -148,12 +180,12 @@ def train(model_type: str, batch_size: int, slide_size: int, img_height: int, im
     callbacks_list = [save_best, early_stop, logs]
 
     model.compile(optimizer=optimizer,
-                  loss=loss, )
+                  loss=loss,)
-    model.fit({'img_in': images}, {'angle_out': angle_cat_array, }, batch_size=batch_size,
+    model.fit({'img_in': images}, {output_name: input_value_array, }, batch_size=batch_size,
               epochs=100, verbose=1, validation_split=0.2, shuffle=True, callbacks=callbacks_list)
 
     # Save model for tensorflow using
-    model.save("/opt/ml/model/tfModel", save_format="tf")
+    model.save(f'/opt/ml/model/model_{record_field.replace("_", "")}_{model_type}_{str(img_width)}x{str(img_height)}h{str(horizon)}')
 
     def representative_dataset() -> typing.Generator[typing.List[tf.float32], typing.Any, None]:
         for d in tf.data.Dataset.from_tensor_slices(images).batch(1).take(100):
@@ -172,9 +204,9 @@ def train(model_type: str, batch_size: int, slide_size: int, img_height: int, im
     tflite_model = converter.convert()
 
     # Save the model.
-    with open(file=f'/opt/ml/model/model_{model_type}_{img_width}x{img_height}h{horizon}.tflite',
+    with open(f'/opt/ml/model/model_{record_field.replace("_", "")}_{model_type}_{str(img_width)}x{str(img_height)}h{str(horizon)}.tflite',
-              mode='wb') as f_output:
+              'wb') as f:
-        f_output.write(tflite_model)
+        f.write(tflite_model)
 
 
 def conv2d(filters: float, kernel: typing.Union[int, typing.Tuple[int, int]], strides: typing.Union[int, typing.Tuple[int, int]], layer_num: int,
@@ -222,20 +254,22 @@ def core_cnn_layers(img_in: Input, img_height: int, img_width: int, drop: float,
     return x
 
 
-def default_linear(input_shape: typing.Tuple[int, int, int] = (120, 160, 3), drop: float = 0.2) -> Model:
+def default_linear(input_shape: typing.Tuple[int, int, int] = (120, 160, 3), drop: float = 0.2,
+                   output_name: str ='angle_out') -> Model:
     img_in = Input(shape=input_shape, name='img_in')
     x = core_cnn_layers(img_in, img_width=input_shape[1], img_height=input_shape[0], drop=drop)
     x = Dense(100, activation='relu', name='dense_1')(x)
     x = Dropout(drop)(x)
     x = Dense(50, activation='relu', name='dense_2')(x)
     x = Dropout(drop)(x)
-    angle_out = Dense(1, activation='linear', name='angle_out')(x)
+    value_out = Dense(1, activation='linear', name=output_name)(x)
 
-    model = Model(inputs=[img_in], outputs=[angle_out], name='linear')
+    model = Model(inputs=[img_in], outputs=[value_out], name='linear')
     return model
 
 
-def default_categorical(input_shape: typing.Tuple[int, int, int] = (120, 160, 3), drop: float = 0.2) -> Model:
+def default_categorical(input_shape: typing.Tuple[int, int, int] = (120, 160, 3), drop: float = 0.2,
+                        output_name: str ='angle_out', output_bin: int = 15) -> Model:
     img_in = Input(shape=input_shape, name='img_in')
     x = core_cnn_layers(img_in, img_width=input_shape[1], img_height=input_shape[0], drop=drop, l4_stride=2)
     x = Dense(100, activation='relu', name="dense_1")(x)
@@ -243,9 +277,9 @@ def default_categorical(input_shape: typing.Tuple[int, int, int] = (120, 160, 3)
     x = Dense(50, activation='relu', name="dense_2")(x)
     x = Dropout(drop)(x)
     # Categorical output of the angle into 15 bins
-    angle_out = Dense(15, activation='softmax', name='angle_out')(x)
+    value_out = Dense(output_bin, activation='softmax', name=output_name)(x)
 
-    model = Model(inputs=[img_in], outputs=[angle_out], name='categorical')
+    model = Model(inputs=[img_in], outputs=[value_out], name='categorical')
     return model
 
 
@@ -260,11 +294,13 @@ def main() -> None:
     parser.add_argument("--batch_size", type=int, default=32)
     parser.add_argument("--drop", type=float, default=0.2)
     parser.add_argument("--model_type", type=str, default=MODEL_CATEGORICAL)
+    parser.add_argument("--record_field", type=str, choices=['angle', 'speed_zone'], default='angle')
 
     args = parser.parse_args()
     params = vars(args)
     train(
         model_type=params["model_type"],
+        record_field=params["record_field"],
         batch_size=params["batch_size"],
         slide_size=params["slide_size"],
         img_height=params["img_height"],