WIP: save work.

python/oddkiva/shakti/inference/darknet/__init__.py

@@ -6,5 +6,6 @@
     RouteSlice,
     RouteConcat,
     Shortcut,
+    Upsample,
     Yolo
 )

python/oddkiva/shakti/inference/darknet/config.py

@@ -207,7 +207,7 @@ def typify_yolo_parameters(self, layer_index):
 
         # The following parameters must be present in the config file.
         self._model[layer_index] = {
-            'upsample': {
+            'yolo': {
                 'mask': mask,
                 'anchors': anchors,
                 'classes': classes,

python/oddkiva/shakti/inference/darknet/network.py

@@ -10,6 +10,14 @@ class Network(nn.Module):
     def __init__(self, cfg: darknet.Config):
         super(Network, self).__init__()
 
+        input_shape = (
+            None,
+             cfg._metadata['channels'],
+             cfg._metadata['height'],
+             cfg._metadata['width']
+        )
+        self.in_shape_at_block = [input_shape]
+        self.out_shape_at_block = [input_shape]
         self.model = self.create_network(cfg)
 
     def create_network(self, cfg: darknet.Config):
@@ -20,60 +28,148 @@ def create_network(self, cfg: darknet.Config):
 
         print(cfg._metadata)
 
-        in_channels_at_block = [cfg._metadata['channels']]
         conv_id = 0
         route_id = 0
         max_pool_id = 0
+        upsample_id = 0
+        yolo_id = 0
 
         for block in cfg._model:
             layer_name = list(block.keys())[0]
             layer_params = block[layer_name]
 
-            print(f'{in_channels_at_block}')
-
             if layer_name == 'convolutional':
-                out_channels = layer_params['filters']
-                stride = layer_params['stride']
-                model.append(darknet.ConvBNA(in_channels_at_block[-1], layer_params,
-                                             conv_id))
-                print(f'[Conv{conv_id}]: Cin {in_channels_at_block[-1]}, Cout {out_channels}, stride {stride}')
-
-                # Update the parameters for the next network block.
-                in_channels_at_block.append(out_channels)
+                self._append_conv(model, layer_params, conv_id)
                 conv_id += 1
             elif layer_name == 'route':
-                layers = layer_params['layers']
-                groups = layer_params['groups']
-                group_id = layer_params['group_id']
-                if len(layers) == 1:
-                    model.append(darknet.RouteSlice(
-                        layers[0], groups, group_id, route_id))
-
-                else:
-                    model.append(darknet.RouteConcat(layers, route_id))
-
-                # Update t
-                if len(layers) == 1:
-                    out_channels = in_channels_at_block[-1] // groups
-                    print(f'[Route{route_id}]: (Slide) Cin = {in_channels_at_block[-1]}, Cout {out_channels}')
-                else:
-                    out_channels = sum([in_channels_at_block[l-1] for l in layers])
-                    print(
-                        f'[Route{route_id}]: '
-                        f'(Concat) Cin={in_channels_at_block[-1]}, Cout={out_channels}, layers={layers}'
-                    )
-
-                in_channels_at_block.append(out_channels)
+                self._append_route(model, layer_params, route_id)
                 route_id += 1
             elif layer_name == 'maxpool':
-                raise NotImplementedError
+                self._append_max_pool(model, layer_params, max_pool_id)
+                max_pool_id += 1
+            elif layer_name == 'upsample':
+                self._append_upsample(model, layer_params, upsample_id)
+                upsample_id += 1
+            elif layer_name == 'yolo':
+                self._append_yolo(model, layer_params, yolo_id)
+                yolo_id += 1
             else:
                 raise NotImplementedError(
                     f'Pytorch layer "{layer_name}" is not implemented'
                 )
 
         return model
 
+    def _append_conv(self, model, layer_params, conv_id):
+        # Extract the input shape.
+        shape_in = self.out_shape_at_block[-1]
+
+        # Calculate the output shape.
+        n, c_in, h_in, w_in = shape_in
+        stride = layer_params['stride']
+        c_out = layer_params['filters']
+        h_out, w_out = h_in // stride, w_in // stride
+        shape_out = (n, c_out, h_out, w_out)
+
+        # Store.
+        self.in_shape_at_block.append(shape_in)
+        self.out_shape_at_block.append(shape_out)
+
+        # Append the convolutional block to the model.
+        model.append(darknet.ConvBNA(c_in, layer_params, conv_id))
+        print(f'[Conv{conv_id}]: '
+              f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
+
+    def _append_route(self, model, layer_params, route_id):
+        layers = layer_params['layers']
+        groups = layer_params['groups']
+        group_id = layer_params['group_id']
+
+        if len(layers) == 1:
+            # Extract the input shape
+            shape_in = self.out_shape_at_block[layers[0]]
+
+            # Calculate the output shape
+            n, c_in, h_in, w_in = shape_in
+            shape_out = (n, c_in // groups, h_in, w_in)
+
+            # Store.
+            self.in_shape_at_block.append(shape_in)
+            self.out_shape_at_block.append(shape_out)
+
+            # Append the route-slice block.
+            model.append(darknet.RouteSlice(
+                layers[0], groups, group_id, route_id))
+            print(f'[Route{route_id}] (Slide): '
+                  f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
+        else:
+            # Fetch all the input shapes.
+            shape_ins = [self.out_shape_at_block[l] for l in layers]
+
+            # Calculate the output shape.
+            n, _, h_in, w_in = shape_ins[0]
+            c_out = sum([shape_in[1] for shape_in in shape_ins])
+            shape_out = (n, c_out, h_in, w_in)
+
+            # Store.
+            self.in_shape_at_block.append(shape_ins)
+            self.out_shape_at_block.append(shape_out)
+
+            # Append the route-concat block.
+            model.append(darknet.RouteConcat(layers, route_id))
+            print(f'[Route{route_id}] (Concat): '
+                  f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
+
+    def _append_max_pool(self, model, layer_params, max_pool_id):
+        # Extract the input shape
+        shape_in = self.out_shape_at_block[-1]
+
+        # Calculate the output shape.
+        size = layer_params['size']
+        stride = layer_params['stride']
+        n, c_in, h_in, w_in = shape_in
+        shape_out = (n, c_in, h_in // stride, w_in // stride)
+
+        # Store.
+        self.in_shape_at_block.append(shape_in)
+        self.out_shape_at_block.append(shape_out)
+
+        model.append(darknet.MaxPool(size, stride))
+        print(f'[MaxPool{max_pool_id}] '
+              f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
+
+    def _append_upsample(self, model, layer_params, upsample_id):
+        # Extract the input shape
+        shape_in = self.out_shape_at_block[-1]
+
+        # Calculate the output shape.
+        stride = layer_params['stride']
+        n, c_in, h_in, w_in = shape_in
+        shape_out = (n, c_in, h_in * stride, w_in * stride)
+
+        # Store.
+        self.in_shape_at_block.append(shape_in)
+        self.out_shape_at_block.append(shape_out)
+
+        model.append(darknet.Upsample(stride))
+        print(f'[Upsample{upsample_id}] '
+              f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
+
+    def _append_yolo(self, model, layer_params, yolo_id):
+        # Extract the input shape
+        shape_in = self.out_shape_at_block[-1]
+
+        # Calculate the output shape.
+        shape_out = shape_in
+
+        # Store.
+        self.in_shape_at_block.append(shape_in)
+        self.out_shape_at_block.append(shape_out)
+
+        model.append(darknet.Yolo(layer_params))
+        print(f'[YOLO{yolo_id}] '
+              f'{self.in_shape_at_block[-1]} -> {self.out_shape_at_block[-1]}')
+
     def load_convolutional_weights(self, conv, weights_file: Path):
         pass
         # with open(weights_file, 'rb') as fp:

python/oddkiva/shakti/inference/darknet/torch_layers.py

@@ -140,13 +140,13 @@ def __init__(self, layers: [int], id: Optional[int] = None):
         self.id = id
 
     def forward(self, x1, x2):
-        if self.layers != 2:
-            raise RuntimeError("Group only 2 inputs")
+        if len(self.layers) != 2:
+            raise RuntimeError("This route-concat layer requires 2 inputs")
         return torch.cat((x1, x2), 1)
 
     def forward(self, x1, x2, x3, x4):
-        if self.layers != 4:
-            raise RuntimeError("Group only 4 inputs")
+        if len(self.layers) != 4:
+            raise RuntimeError("This route-concat layer requires 4 inputs")
         return torch.cat((x1, x2, x3, x4), 1)
 
 
@@ -161,7 +161,9 @@ def __init__(self, activation: str):
         elif activation == 'relu':
             self.activation_fn = ReLU(inplace=True)
         else:
-            raise NotImplementedError('Activation function not implemented!')
+            raise NotImplementedError(
+                f'The followig activation function "{activation}" not implemented!'
+            )
 
     def forward(self, x1, x2):
         x = self.activation_fn(x1 + x2)
@@ -180,9 +182,34 @@ def forward(self, x):
 
 class Yolo(nn.Module):
 
-    def __init__(self):
+    def __init__(self, darknet_params: dict[str, Any]):
         super(Yolo, self).__init__()
+        self.masks = darknet_params['mask']
+        self.anchors = darknet_params['anchors']
+        self.scale_x_y = darknet_params['scale_x_y']
+        self.classes = darknet_params['classes']
+
+        self.alpha = self.scale_x_y
+        self.beta = -0.5 * (self.scale_x_y - 1)
 
     def forward(self, x):
-        raise NotImplementedError()
-        return x;
+        num_box_features = 5 + self.num_classes
+        assert num_box_features == 85
+
+        y = x
+
+        # Box prediction
+        xs = [box * num_box_features + 0 for box in range(3)]
+        ys = [box * num_box_features + 1 for box in range(3)]
+        # ws = [box * num_box_features + 2 for box in range(3)]
+        # hs = [box * num_box_features + 3 for box in range(3)]
+        y[:, xs, :, :] = self.alpha * nn.Sigmoid(x[:, xs, :, :]) + self.beta
+        y[:, ys, :, :] = self.alpha * nn.Sigmoid(x[:, ys, :, :]) + self.beta
+
+        # Class probabilities.
+        for box in range(0, 3):
+            c_begin = box * num_box_features + 4
+            c_end = (box + 1) * num_box_features 
+            y[:, c_begin:c_end, :, :] = nn.Sigmoid(x[:, c_begin:c_end, :, :])
+
+        return y;