Adding the main SyNET branches

requirements.txt

@@ -0,0 +1,9 @@
+gradio
+ultralytics
+transformers
+datasets
+trl
+peft
+tf-keras
+bitsandbytes
+pytest

synet/__init__.py

@@ -0,0 +1,15 @@
+from .backends import get_backend
+
+
+__all__ = "backends", "base", "katana", "sabre", "quantize", "test", \
+    "metrics", "tflite_utils"
+
+
+def get_model(model_path, backend, *args, **kwds):
+    """Method to get the model.  For now, only the katananet model is
+supported in ultralytics format."""
+
+    print("loading", model_path)
+
+    backend = get_backend(backend)
+    return backend.get_model(model_path, *args, **kwds)

synet/__main__.py

@@ -0,0 +1,14 @@
+from importlib import import_module
+from sys import argv, exit
+
+import synet
+
+
+def main():
+    if argv[1] in synet.__all__:
+        return import_module(f"synet.{argv.pop(1)}").main()
+    return import_module(f"synet.backends.{argv.pop(1)}").main()
+
+
+if __name__ == "__main__":
+    exit(main())

synet/asymmetric.py

@@ -0,0 +1,113 @@
+"""asymetric.py diverges from base.py and layers.py in that its core
+assumption is switched.  In base.py/layers.py, the output of a module
+in keras vs torch is identical, while asymetric modules act as the
+identity function in keras.  To get non-identity behavior in keras
+mode, you should call module.clf().  'clf' should be read as 'channels
+last forward'; such methods take in and return a channels-last numpy
+array.
+
+The main use case for these modules is for uniform preprocessing to
+bridge the gap between 'standard' training scenarios and actual
+execution environments.  So far, the main examples implemented are
+conversions to grayscale, bayer, and camera augmented images.  This
+way, you can train your model on a standard RGB pipeline.  The
+resulting tflite model will not have these extra layers, and is ready
+to operate on the raw input at deployment.
+
+The cfl methods are mainly used for python demos where the sensor
+still needs to be simulated, but not included in the model.
+
+"""
+
+from os.path import join, dirname
+from cv2 import GaussianBlur as cv2GaussianBlur
+from numpy import array, interp, ndarray
+from numpy.random import normal
+from torch import empty, tensor, no_grad, rand
+from torchvision.transforms import GaussianBlur
+
+from .demosaic import Demosaic, UnfoldedDemosaic, Mosaic
+from .base import askeras, Module
+
+
+class Grayscale(Module):
+    """Training frameworks often fix input channels to 3.  This
+grayscale layer can be added to the beginning of a model to convert to
+grayscale.  This layer is ignored when converting to tflite.  The end
+result is that the pytorch model can take any number of input
+channels, but the tensorflow (tflite) model expects exactly one input
+channel.
+
+    """
+
+    def forward(self, x):
+        if askeras.use_keras:
+            return x
+        return x.mean(1, keepdims=True)
+
+
+class Camera(Module):
+    def __init__(self,
+                 gamma,
+                 bayer_pattern='gbrg',
+                 from_bayer=False,
+                 to_bayer=False,
+                 ratio=(1, 1, 1),
+                 blur_sigma=0.4,
+                 noise_sigma=10/255):
+        super().__init__()
+        self.mosaic = Mosaic(bayer_pattern)
+        self.demosaic = UnfoldedDemosaic('malvar', bayer_pattern
+                                         ).requires_grad_(False)
+        self.blur_sigma = blur_sigma
+        self.noise_sigma = noise_sigma
+        self.from_bayer = from_bayer
+        self.to_bayer = to_bayer
+        self.gamma = gamma
+        self.blur = GaussianBlur(3, blur_sigma)
+        self.ratio = ratio
+
+    def gamma_correction(self, image):
+
+        for yoff, xoff, chan in zip(self.mosaic.rows,
+                                    self.mosaic.cols,
+                                    self.mosaic.bayer_pattern):
+            # the gamma correction (from experiments) is channel dependent
+            image[yoff::2, xoff::2] = ((image[yoff::2, xoff::2]
+                                        ) ** self.gamma[chan])
+        return image
+
+    #@no_grad
+    def forward(self, im):
+        if askeras.use_keras:
+            return im
+        if not self.from_bayer:
+            im = self.mosaic(im)
+        if rand(1) < self.ratio[0]:
+            im = self.blur(im)
+        if rand(1) < self.ratio[1]:
+            im = self.gamma_correction(im)
+        if rand(1) < self.ratio[2]:
+            this_noise_sigma, = empty(1).normal_(self.noise_sigma, 2/255)
+            im += empty(im.shape, device=im.device
+                        ).normal_(0.0, max(0, this_noise_sigma))
+        if not self.to_bayer:
+            im = self.demosaic(im)
+        return im.clip(0, 1)
+
+    def clf(self, im):
+        assert False, "didn't update this function after refactor"
+        # augmentation should always be done on bayer image.
+        if not self.from_bayer:
+            im = self.mosaic.clf(im)
+        # let the noise level vary
+        this_noise_sigma = normal(self.noise_sigma, 2)
+        # if you blur too much, the image becomes grayscale
+        im = cv2GaussianBlur(im, [3, 3], self.blur_sigma)
+        im = self.map_to_linear(im)
+        # GaussianBlur likes to remove singleton channel dimension
+        im = im[..., None] + normal(0.0, this_noise_sigma, im.shape + (1,))
+        # depending on scenario, you may not want to return an RGB image.
+        if not self.to_bayer:
+            im = self.demosaic.clf(im)
+        return im.clip(0, 255)

synet/backends/__init__.py

@@ -0,0 +1,68 @@
+from importlib import import_module
+from shutil import copy
+
+from ..zoo import in_zoo, get_config, get_configs, get_weights
+
+
+class Backend:
+    def get_model(self, model_path):
+        """Load model from config, or pretrained save."""
+        raise NotImplementedError("Please subclass and implement")
+
+    def get_shape(self, model):
+        """Get shape of model."""
+        raise NotImplementedError("Please subclass and implement")
+
+    def patch(self):
+        """Initialize backend to utilize Synet Modules"""
+        raise NotImplementedError("Please subclass and implement")
+
+    def val_post(self, weights, tflite, val_post, conf_thresh=.25,
+                 iou_thresh=.7):
+        """Default conf_thresh and iou_thresh (.25 and .75 resp.)
+        taken from ultralytics/cfg/default.yaml.
+
+        """
+        raise NotImplementedError("Please subclass and implement")
+
+    def tf_post(self, tflite, val_post, conf_thresh, iou_thresh):
+        """Loads the tflite, loads the image, preprocesses the image,
+        evaluates the tflite on the pre-processed image, and performs
+        post-processing on the tflite output with a given confidence
+        and iou threshold.
+
+        :param tflite: Path to tflite file, or a raw tflite buffer
+        :param val_post: Path to image to evaluate on.
+        :param conf_thresh: Confidence threshould.  See val_post docstring
+        above for default value details.
+        :param iou_thresh: IoU threshold for NMS.  See val_post docstring
+        above for default value details.
+
+        """
+        raise NotImplementedError("Please subclass and implement")
+
+    def get_chip(self, model):
+        """Get chip of model."""
+        raise NotImplementedError("Please subclass and implement")
+
+    def maybe_grab_from_zoo(self, model_path):
+        if in_zoo(model_path, self.name):
+            copy(get_config(model_path, self.name), model_path)
+        elif model_path.endswith(".pt") or model_path.endswith(".tflite"):
+            get_weights(model_path, self.name)
+        return model_path
+
+    def get_configs(self):
+        return get_configs(self.name)
+
+    def get_data(self, data):
+        """return a {split:files} where files is either a string or
+        list of strings denoting path(s) to file(s) or
+        directory(ies). Files should be newline-seperated lists of
+        image paths, and directories should (recursively) contain only
+        images or directories."""
+        raise NotImplementedError("Please subclass and implement")
+
+
+def get_backend(name):
+    return import_module(f".{name}", __name__).Backend()

synet/backends/custom.py

@@ -0,0 +1,82 @@
+from .base import askeras
+
+from object_detection.models.tf import TFDetect as PC_TFDetect
+from tensorflow.math import ceil
+import tensorflow as tf
+class TFDetect(PC_TFDetect):
+    def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None):
+        super().__init__(nc, anchors, ch, imgsz, w)
+        for i in range(self.nl):
+            ny, nx = (ceil(self.imgsz[0] / self.stride[i]),
+                      ceil(self.imgsz[1] / self.stride[i]))
+            self.grid[i] = self._make_grid(nx, ny)
+
+    # copy call method, but replace // with ceil div
+    def call(self, inputs):
+        if askeras.kwds.get('deploy'):
+            return self.deploy(inputs)
+        z = []  # inference output
+        x = []
+        for i in range(self.nl):
+            x.append(self.m[i](inputs[i]))
+            # x(bs,20,20,255) to x(bs,3,20,20,85)
+            ny, nx = (ceil(self.imgsz[0] / self.stride[i]),
+                      ceil(self.imgsz[1] / self.stride[i]))
+            x[i] = tf.transpose(tf.reshape(x[i], [-1, ny * nx, self.na, self.no]), [0, 2, 1, 3])
+
+            if not self.training:  # inference
+                y = tf.sigmoid(x[i])
+                xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i]  # xy
+                wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]
+                # Normalize xywh to 0-1 to reduce calibration error
+                xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
+                wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
+                y = tf.concat([xy, wh, y[..., 4:]], -1)
+                # y = tf.concat([xy, wh, y[..., 4:]], 3)
+                z.append(tf.reshape(y, [-1, self.na * ny * nx, self.no]))
+
+        return x if self.training else (tf.concat(z, 1), x)
+
+    def deploy(self, inputs):
+        assert inputs[0].shape[0] == 1, 'requires batch_size == 1'
+        box1, box2, cls = [], [], []
+        for mi, xi, gi, ai, si in zip(self.m, inputs, self.grid, self.anchor_grid, self.stride):
+            x = tf.reshape(tf.sigmoid(mi(xi)), (1, -1, self.na, self.no))
+            xy = (x[..., 0:2] * 2 + (tf.transpose(gi, (0, 2, 1, 3)) - .5)) * si
+            wh = (x[..., 2:4] * 2) ** 2 * tf.transpose(ai, (0, 2, 1, 3))
+            box1.append(tf.reshape(xy - wh/2, (1, -1, 2)))
+            box2.append(tf.reshape(xy + wh/2, (1, -1, 2)))
+            cls.append(tf.reshape(x[..., 4:5]*x[..., 5:], (1, -1, x.shape[-1]-5)))
+        return (tf.concat(box1, 1, name='box1'),
+                tf.concat(box2, 1, name='box2'),
+                tf.concat(cls,  1, name='cls'))
+
+
+from object_detection.models.yolo import Detect as PC_PTDetect
+class Detect(PC_PTDetect):
+    def __init__(self, *args, **kwds):
+        if len(args) == 4:
+            args = args[:3]
+        # construct normally
+        super().__init__(*args, **kwds)
+        # save args/kwargs for later construction of TF model
+        self.args = args
+        self.kwds = kwds
+    def forward(self, x, theta=None):
+        if askeras.use_keras:
+            assert theta is None
+            return self.as_keras(x)
+        return super().forward(x, theta=theta)
+    def as_keras(self, x):
+        return TFDetect(*self.args, imgsz=askeras.kwds["imgsz"],
+                        w=self, **self.kwds
+                        )(x)
+
+from object_detection.models import yolo
+from importlib import import_module
+def patch_custom(chip):
+    # patch custom.models.yolo
+    module = import_module(f'..{chip}', __name__)
+    setattr(yolo, chip, module)
+    yolo.Concat = module.Cat
+    yolo.Detect = module.Detect = Detect

synet/backends/ultralytics.py

@@ -0,0 +1,404 @@
+
+from importlib import import_module
+from sys import argv
+
+from cv2 import imread, imwrite, resize
+from numpy import array
+from torch import tensor
+from torch.nn import ModuleList
+from ultralytics import YOLO
+from ultralytics.data.utils import check_det_dataset, check_cls_dataset
+from ultralytics.engine import validator, predictor, trainer
+from ultralytics.engine.results import Results
+from ultralytics.models.yolo import model as yolo_model
+from ultralytics.nn import tasks
+from ultralytics.nn.autobackend import AutoBackend
+from ultralytics.nn.modules.block import DFL as Torch_DFL, Proto as Torch_Proto
+from ultralytics.nn.modules.head import (Pose as Torch_Pose,
+                                         Detect as Torch_Detect,
+                                         Segment as Torch_Segment,
+                                         Classify as Torch_Classify)
+from ultralytics.utils import dist
+from ultralytics.utils.ops import non_max_suppression, process_mask
+from ultralytics.utils.checks import check_imgsz
+
+from . import Backend as BaseBackend
+from ..base import (askeras, Conv2d, ReLU, Upsample, GlobalAvgPool,
+                    Dropout, Linear)
+from .. import layers
+from .. import asymmetric
+from ..layers import Sequential, CoBNRLU
+from ..tflite_utils import tf_run, concat_reshape
+
+
+class DFL(Torch_DFL):
+    def __init__(self, c1=16, sm_split=None):
+        super().__init__(c1)
+        weight = self.conv.weight
+        self.conv = Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
+        self.conv.conv.weight.data[:] = weight.data
+        self.sm_split = sm_split
+
+    def forward(self, x):
+        if askeras.use_keras:
+            return self.as_keras(x)
+        return super().forward(x)
+
+    def as_keras(self, x):
+        # b, ay, ax, c = x.shape
+        from tensorflow.keras.layers import Reshape, Softmax
+        if hasattr(self, "sm_split") and self.sm_split is not None:
+            from tensorflow.keras.layers import Concatenate
+            assert not (x.shape[0]*x.shape[1]*x.shape[2]*4) % self.sm_split
+            x = Reshape((self.sm_split, -1, self.c1))(x)
+            # tensorflow really wants to be indented like this.  I relent...
+            return Reshape((-1, 4))(
+                self.conv(
+                    Concatenate(1)([
+                        Softmax(-1)(x[:, i:i+1])
+                        for i in range(x.shape[1])
+                    ])
+                )
+            )
+
+        return Reshape((-1, 4)
+                       )(self.conv(Softmax(-1)(Reshape((-1, 4, self.c1))(x))))
+
+
+class Proto(Torch_Proto):
+    def __init__(self, c1, c_=256, c2=32):
+        """arguments understood as in_channels, number of protos, and
+        number of masks"""
+        super().__init__(c1, c_, c2)
+        self.cv1 = CoBNRLU(c1, c_, 3)
+        self.upsample = Upsample(scale_factor=2, mode='bilinear')
+        self.cv2 = CoBNRLU(c_, c_, 3)
+        self.cv3 = CoBNRLU(c_, c2, 1, name='proto')
+
+
+def generate_anchors(H, W, stride, offset):
+    from tensorflow import meshgrid, range, stack, reshape, concat
+    from tensorflow.math import ceil
+    return concat([stack((reshape((sx + offset) * s, (-1,)),
+                          reshape((sy + offset) * s, (-1,))),
+                         -1)
+                   for s, (sy, sx) in ((s.item(),
+                                        meshgrid(range(ceil(H/s)),
+                                                 range(ceil(W/s)),
+                                                 indexing="ij"))
+                                       for s in stride)],
+                  -2)
+
+
+class Detect(Torch_Detect):
+    def __init__(self, nc=80, ch=(), sm_split=None, junk=None):
+        super().__init__(nc, ch)
+        c2 = max((16, ch[0] // 4, self.reg_max * 4))
+        self.cv2 = ModuleList(Sequential(Conv2d(x, c2, 3, bias=True),
+                                         ReLU(6),
+                                         Conv2d(c2, c2, 3, bias=True),
+                                         ReLU(6),
+                                         Conv2d(c2, 4 * self.reg_max, 1,
+                                                bias=True))
+                              for x in ch)
+        self.cv3 = ModuleList(Sequential(Conv2d(x, x, 3, bias=True),
+                                         ReLU(6),
+                                         Conv2d(x, x, 3, bias=True),
+                                         ReLU(6),
+                                         Conv2d(x, self.nc, 1, bias=True))
+                              for x in ch)
+        if junk is None:
+            sm_split = None
+        self.dfl = DFL(sm_split=sm_split)
+
+    def forward(self, x):
+        if askeras.use_keras:
+            return Detect.as_keras(self, x)
+        return super().forward(x)
+
+    def as_keras(self, x):
+        from tensorflow.keras.layers import Reshape
+        from tensorflow import stack
+        from tensorflow.keras.layers import (Concatenate, Subtract,
+                                             Add, Activation)
+        from tensorflow.keras.activations import sigmoid
+        ltrb = Concatenate(-2)([self.dfl(cv2(xi)) * s.item()
+                                for cv2, xi, s in
+                                zip(self.cv2, x, self.stride)])
+        H, W = askeras.kwds['imgsz']
+        anchors = generate_anchors(H, W, self.stride, .5)             # Nx2
+        anchors = stack([anchors for batch in range(x[0].shape[0])])  # BxNx2
+        box1 = Subtract(name="box1")((anchors, ltrb[:, :, :2]))
+        box2 = Add(name="box2")((anchors, ltrb[:, :, 2:]))
+        if askeras.kwds.get("xywh"):
+            box1, box2 = (box1 + box2) / 2, box2 - box1
+
+        cls = Activation(sigmoid, name='cls')(
+            Concatenate(-2)([
+                Reshape((-1, self.nc))(cv3(xi))
+                for cv3, xi in zip(self.cv3, x)
+            ])
+        )
+        out = [box1, box2, cls]
+        if askeras.kwds.get("quant_export"):
+            return out
+        # everything after here needs to be implemented by post-processing
+        out[:2] = (box/array((W, H)) for box in out[:2])
+        return Concatenate(-1)(out)
+
+
+class Pose(Torch_Pose, Detect):
+    def __init__(self, nc, kpt_shape, ch, sm_split=None, junk=None):
+        super().__init__(nc, kpt_shape, ch)
+        Detect.__init__(self, nc, ch, sm_split, junk=junk)
+        self.detect = Detect.forward
+        c4 = max(ch[0] // 4, self.nk)
+        self.cv4 = ModuleList(Sequential(Conv2d(x, c4, 3),
+                                         ReLU(6),
+                                         Conv2d(c4, c4, 3),
+                                         ReLU(6),
+                                         Conv2d(c4, self.nk, 1))
+                              for x in ch)
+
+    def forward(self, *args, **kwds):
+        if askeras.use_keras:
+            return self.as_keras(*args, **kwds)
+        return super().forward(*args, **kwds)
+
+    def s(self, stride):
+        if self.kpt_shape[1] == 3:
+            from tensorflow import constant
+            return constant([stride, stride, 1]*self.kpt_shape[0])
+        return stride
+
+    def as_keras(self, x):
+
+        from tensorflow.keras.layers import Reshape, Concatenate, Add
+        from tensorflow import stack, reshape
+        from tensorflow.keras.activations import sigmoid
+
+        if self.kpt_shape[1] == 3:
+            presence_chans = [i*3+2 for i in range(17)]
+            pres, kpts = zip(*((Reshape((-1, self.kpt_shape[0], 1)
+                                        )(presence(xi)),
+                                Reshape((-1, self.kpt_shape[0], 2)
+                                        )(keypoint(xi)*s*2))
+                               for presence, keypoint, xi, s in
+                               ((*cv[-1].split_channels(presence_chans),
+                                 cv[:-1](xi), s.item())
+                                for cv, xi, s in
+                                zip(self.cv4, x, self.stride))))
+            pres = Concatenate(-3, name="pres")([sigmoid(p) for p in pres])
+        else:
+            kpts = [Reshape((-1, self.kpt_shape[0], 2))(cv(xi)*s*2)
+                    for cv, xi, s in
+                    zip(self.cv4, x, self.stride)]
+
+        H, W = askeras.kwds['imgsz']
+        anchors = generate_anchors(H, W, self.stride, offset=0)       # Nx2
+        anchors = reshape(anchors, (-1, 1, 2))                        # Nx1x2
+        anchors = stack([anchors for batch in range(x[0].shape[0])])  # BxNx1x2
+        kpts = Add(name='kpts')((Concatenate(-3)(kpts), anchors))
+
+        x = self.detect(self, x)
+
+        if askeras.kwds.get("quant_export"):
+            if self.kpt_shape[1] == 3:
+                return *x, kpts, pres
+            return *x, kpts
+
+        # everything after here needs to be implemented by post-processing
+        if self.kpt_shape[1] == 3:
+            kpts = Concatenate(-1)((kpts, pres))
+
+        return Concatenate(-1)((x, Reshape((-1, self.nk))(kpts)))
+
+
+class Segment(Torch_Segment, Detect):
+    """YOLOv8 Segment head for segmentation models."""
+
+    def __init__(self, nc=80, nm=32, npr=256, ch=(), sm_split=None, junk=None):
+        super().__init__(nc, nm, npr, ch)
+        Detect.__init__(self, nc, ch, sm_split, junk=junk)
+        self.detect = Detect.forward
+        self.proto = Proto(ch[0], self.npr, self.nm)  # protos
+        c4 = max(ch[0] // 4, self.nm)
+        self.cv4 = ModuleList(Sequential(CoBNRLU(x, c4, 3),
+                                         CoBNRLU(c4, c4, 3),
+                                         Conv2d(c4, self.nm, 1))
+                              for x in ch)
+
+    def forward(self, x):
+        if askeras.use_keras:
+            return self.as_keras(x)
+        return super().forward(x)
+
+    def as_keras(self, x):
+        from tensorflow.keras.layers import Reshape, Concatenate
+        p = self.proto(x[0])
+        mc = Concatenate(-2, name='seg')([Reshape((-1, self.nm))(cv4(xi))
+                                          for cv4, xi in zip(self.cv4, x)])
+        x = self.detect(self, x)
+        if askeras.kwds.get("quant_export"):
+            return *x, mc, p
+        # everything after here needs to be implemented by post-processing
+        return Concatenate(-1)((x, mc)), p
+
+
+class Classify(Torch_Classify):
+    def __init__(self, junk, c1, c2, k=1, s=1, p=None, g=1):
+        super().__init__(c1, c2, k=k, s=s, p=p, g=g)
+        c_ = 1280
+        assert p is None
+        self.conv = CoBNRLU(c1, c_, k, s, groups=g)
+        self.pool = GlobalAvgPool()
+        self.drop = Dropout(p=0.0, inplace=True)
+        self.linear = Linear(c_, c2)
+
+    def forward(self, x):
+        if askeras.use_keras:
+            return self.as_keras(x)
+        return super().forward(x)
+
+    def as_keras(self, x):
+        from keras.layers import Concatenate, Flatten, Softmax
+        if isinstance(x, list):
+            x = Concatenate(-1)(x)
+        x = self.linear(self.drop(Flatten()(self.pool(self.conv(x)))))
+        return x if self.training else Softmax()(x)
+
+
+class Backend(BaseBackend):
+
+    models = {}
+    name = "ultralytics"
+
+    def get_model(self, model_path, full=False):
+
+        model_path = self.maybe_grab_from_zoo(model_path)
+
+        if model_path in self.models:
+            model = self.models[model_path]
+        else:
+            model = self.models[model_path] = YOLO(model_path)
+
+        if full:
+            return model
+        return model.model
+
+    def get_shape(self, model):
+        if isinstance(model, str):
+            model = self.get_model(model)
+        return model.yaml["image_shape"]
+
+    def patch(self, model_path=None):
+        for module in layers, asymmetric:
+            for name in dir(module):
+                if name[0] != "_":
+                    setattr(tasks, name, getattr(module, name))
+        tasks.Concat = layers.Cat
+        tasks.Pose = Pose
+        tasks.Detect = Detect
+        tasks.Segment = Segment
+        tasks.Classify = Classify
+        orig_ddp_file = dist.generate_ddp_file
+
+        def generate_ddp_file(trainer):
+            fname = orig_ddp_file(trainer)
+            fstr = open(fname).read()
+            open(fname, 'w').write(f"""\
+from synet.backends import get_backend
+get_backend('ultralytics').patch()
+{fstr}""")
+            return fname
+        dist.generate_ddp_file = generate_ddp_file
+
+        def tflite_check_imgsz(*args, **kwds):
+            kwds['stride'] = 1
+            return check_imgsz(*args, **kwds)
+        trainer.check_imgsz = tflite_check_imgsz
+        if model_path is not None and model_path.endswith('tflite'):
+            print('SyNet: model provided is tflite.  Modifying validators'
+                  ' to anticipate tflite output')
+            task_map = yolo_model.YOLO(model_path).task_map
+            for task in task_map:
+                for mode in 'predictor', 'validator':
+                    class Wrap(task_map[task][mode]):
+                        def postprocess(self, preds, *args, **kwds):
+                            # concate_reshape currently expect ndarry
+                            # with batch size of 1, so remove and
+                            # re-add batch and tensorship.
+                            preds = concat_reshape([p[0].numpy()
+                                                    for p in preds],
+                                                   self.args.task,
+                                                   classes_to_index=False,
+                                                   xywh=True)
+                            if isinstance(preds, tuple):
+                                preds = (tensor(preds[0][None])
+                                         .permute(0, 2, 1),
+                                         tensor(preds[1][None])
+                                         .permute(0, 2, 3, 1))
+                            else:
+                                preds = tensor(preds[None]).permute(0, 2, 1)
+                            return super().postprocess(preds, *args, **kwds)
+                    if task != 'classify':
+                        task_map[task][mode] = Wrap
+            yolo_model.YOLO.task_map = task_map
+
+            class TfliteAutoBackend(AutoBackend):
+                def __init__(self, *args, **kwds):
+                    super().__init__(*args, **kwds)
+                    self.output_details.sort(key=lambda x: x['name'])
+                    if len(self.output_details) == 1:  # classify
+                        num_classes = self.output_details[0]['shape'][-1]
+                    else:
+                        num_classes = self.output_details[2]['shape'][2]
+                    self.kpt_shape = (self.output_details[-1]['shape'][-2], 3)
+                    self.names = {k: self.names[k] for k in range(num_classes)}
+
+            validator.check_imgsz = tflite_check_imgsz
+            predictor.check_imgsz = tflite_check_imgsz
+            validator.AutoBackend = TfliteAutoBackend
+            predictor.AutoBackend = TfliteAutoBackend
+
+    def get_data(self, data):
+        try:
+            return check_det_dataset(data)
+        except Exception as e:
+            try:
+                return check_cls_dataset(data)
+            except Exception as e2:
+                print("unable to load data as classification or detection dataset")
+                print(e2)
+                raise e
+
+
+def main():
+
+    backend = Backend()
+
+    # copy model from zoo if necessary
+    for ind, val in enumerate(argv):
+        if val.startswith("model="):
+            model = backend.maybe_grab_from_zoo(val.split("=")[1])
+            argv[ind] = "model="+model
+
+            # add synet ml modules to ultralytics
+            backend.patch(model_path=model)
+
+            # add imgsz if not explicitly given
+            for val in argv:
+                if val.startswith("imgsz="):
+                    break
+            else:
+                argv.append(f"imgsz={max(backend.get_shape(model))}")
+
+            break
+
+
+    # launch ultralytics
+    try:
+        from ultralytics.cfg import entrypoint
+    except:
+        from ultralytics.yolo.cfg import entrypoint
+    entrypoint()

synet/backends/yolov5.py

@@ -0,0 +1,436 @@
+from types import SimpleNamespace
+from importlib import import_module
+
+import numpy
+import tensorflow as tf
+from tensorflow.math import ceil
+from torch import load, no_grad, tensor
+from yolov5 import val
+from yolov5.models import yolo, common
+from yolov5.models.yolo import Detect as Yolo_PTDetect, Model
+from yolov5.models.tf import TFDetect as Yolo_TFDetect
+from yolov5.utils.general import non_max_suppression
+from yolov5.val import (Path, Callbacks, create_dataloader,
+                        select_device, DetectMultiBackend,
+                        check_img_size, LOGGER, check_dataset, torch,
+                        np, ConfusionMatrix, coco80_to_coco91_class,
+                        Profile, tqdm, scale_boxes, xywh2xyxy,
+                        output_to_target, ap_per_class, pd,
+                        increment_path, os, colorstr, TQDM_BAR_FORMAT,
+                        process_batch, plot_images, save_one_txt)
+
+from .base import askeras
+
+
+def get_yolov5_model(model_path, low_thld=0, raw=False, **kwds):
+    """Convenience function to load yolov5 model"""
+    if model_path.endswith(".yml") or model_path.endswith(".yaml"):
+        assert raw
+        return Model(model_path)
+    ckpt = load(model_path)
+    ckpt = ckpt['model'] if isinstance(ckpt, dict) else ckpt
+    raw_model = Model(ckpt.yaml)
+    raw_model.load_state_dict(ckpt.state_dict())
+    if raw:
+        return raw_model
+    raw_model.eval()
+
+    def model(x):
+        with no_grad():
+            xyxyoc = non_max_suppression(raw_model(tensor(x)
+                                                   .unsqueeze(0).float()),
+                                         conf_thres=low_thld,
+                                         iou_thres=.3,
+                                         multi_label=True
+                                         )[0].numpy()
+            return xyxyoc[:, :4], xyxyoc[:, 4:].prod(1)
+    return model
+
+
+class TFDetect(Yolo_TFDetect):
+    """Modify Tensorflow Detect head to allow for arbitrary input
+    shape (need not be multiple of 32).
+
+    """
+
+    # use orig __init__, but make nx, ny calculated via ceil div
+    def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None):
+        super().__init__(nc, anchors, ch, imgsz, w)
+        for i in range(self.nl):
+            ny, nx = (ceil(self.imgsz[0] / self.stride[i]),
+                      ceil(self.imgsz[1] / self.stride[i]))
+            self.grid[i] = self._make_grid(nx, ny)
+
+    # copy call method, but replace // with ceil div
+    def call(self, inputs):
+        z = []  # inference output
+        x = []
+        for i in range(self.nl):
+            x.append(self.m[i](inputs[i]))
+            # x(bs,20,20,255) to x(bs,3,20,20,85)
+            ny, nx = (ceil(self.imgsz[0] / self.stride[i]),
+                      ceil(self.imgsz[1] / self.stride[i]))
+            x[i] = tf.reshape(x[i], [-1, ny * nx, self.na, self.no])
+
+            if not self.training:  # inference
+                y = x[i]
+                grid = tf.transpose(self.grid[i], [0, 2, 1, 3]) - 0.5
+                anchor_grid = tf.transpose(self.anchor_grid[i], [0, 2, 1, 3])*4
+                xy = (tf.sigmoid(y[..., 0:2]) * 2 + grid) * self.stride[i]
+                wh = tf.sigmoid(y[..., 2:4]) ** 2 * anchor_grid
+                # Normalize xywh to 0-1 to reduce calibration error
+                xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]],
+                                  dtype=tf.float32)
+                wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]],
+                                  dtype=tf.float32)
+                y = tf.concat([xy, wh, tf.sigmoid(y[..., 4:5 + self.nc]),
+                               y[..., 5 + self.nc:]], -1)
+                z.append(tf.reshape(y, [-1, self.na * ny * nx, self.no]))
+
+        return tf.transpose(x, [0, 2, 1, 3]) \
+            if self.training else (tf.concat(z, 1),)
+
+
+class Detect(Yolo_PTDetect):
+    """Make YOLOv5 Detect head compatible with synet tflite export"""
+
+    def __init__(self, *args, **kwds):
+        # to account for args hack.
+        if len(args) == 4:
+            args = args[:3]
+        # construct normally
+        super().__init__(*args, **kwds)
+        # save args/kwargs for later construction of TF model
+        self.args = args
+        self.kwds = kwds
+
+    def forward(self, x):
+        if askeras.use_keras:
+            return self.as_keras(x)
+        return super().forward(x)
+
+    def as_keras(self, x):
+        return TFDetect(*self.args, imgsz=askeras.kwds["imgsz"],
+                        w=self, **self.kwds
+                        )(x)
+
+
+def val_run_tflite(
+        data,
+        weights=None,  # model.pt path(s)
+        batch_size=None,  # batch size
+        batch=None,  # batch size
+        imgsz=None,  # inference size (pixels)
+        img=None,  # inference size (pixels)
+        conf_thres=0.001,  # confidence threshold
+        iou_thres=0.6,  # NMS IoU threshold
+        max_det=300,  # maximum detections per image
+        task='val',  # train, val, test, speed or study
+        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+        workers=8,  # max dataloader workers (per RANK in DDP mode)
+        single_cls=False,  # treat as single-class dataset
+        augment=False,  # augmented inference
+        verbose=False,  # verbose output
+        save_txt=False,  # save results to *.txt
+        save_hybrid=False,  # save label+prediction hybrid results to *.txt
+        save_conf=False,  # save confidences in --save-txt labels
+        save_json=False,  # save a COCO-JSON results file
+        project='runs/val',  # save to project/name
+        name='exp',  # save to project/name
+        exist_ok=False,  # existing project/name ok, do not increment
+        half=True,  # use FP16 half-precision inference
+        dnn=False,  # use OpenCV DNN for ONNX inference
+        model=None,
+        dataloader=None,
+        save_dir=Path(''),
+        plots=True,
+        callbacks=Callbacks(),
+        compute_loss=None,
+):
+
+    if imgsz is None and img is None:
+        imgsz = 640
+    elif img is not None:
+        imgsz = img
+    if batch_size is None and batch is None:
+        batch_size = 32
+    elif batch is not None:
+        batch_size = batch
+
+    # Initialize/load model and set device
+    training = model is not None
+    if training:  # called by train.py
+        device, pt, jit, engine = next(model.parameters()).device, True, False, False  # get model device, PyTorch model
+        half &= device.type != 'cpu'  # half precision only supported on CUDA
+        model.half() if half else model.float()
+
+        # SYNET MODIFICATION: never train tflite
+        tflite = False
+
+    else:  # called directly
+        device = select_device(device, batch_size=batch_size)
+        half &= device.type != 'cpu' # half precision only supported on CUDA, dont remove!
+
+        # Directories
+        save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
+        (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+
+        # Load model
+        model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
+
+        # SYNET MODIFICATION: check for tflite
+        tflite = hasattr(model, "interpreter")
+
+        stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
+
+        # SYNET MODIFICATION: if tflite, use that shape
+        if tflite:
+            sn = model.input_details[0]['shape']
+            imgsz = int(max(sn[2], sn[1]))
+
+        if not isinstance(imgsz, (list, tuple)):
+            imgsz = check_img_size(imgsz, s=stride)  # check image size
+        half = model.fp16  # FP16 supported on limited backends with CUDA
+        if engine:
+            batch_size = model.batch_size
+        else:
+            device = model.device
+            if not (pt or jit):
+                batch_size = 1  # export.py models default to batch-size 1
+                LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')
+
+        # Data
+        data = check_dataset(data)  # check
+
+    # Configure
+    model.eval()
+    cuda = device.type != 'cpu' # half precision only supported on CUDA, dont remove!
+    is_coco = isinstance(data.get('val'), str) and data['val'].endswith(f'coco{os.sep}val2017.txt')  # COCO dataset
+    nc = 1 if single_cls else int(data['nc'])  # number of classes
+    iouv = torch.linspace(0.5, 0.95, 10, device=device)  # iou vector for mAP@0.5:0.95
+    niou = iouv.numel()
+
+    # Dataloader
+    if not training:
+        if pt and not single_cls:  # check --weights are trained on --data
+            ncm = model.model.nc
+            assert ncm == nc, f'{weights} ({ncm} classes) trained on different --data than what you passed ({nc} ' \
+                              f'classes). Pass correct combination of --weights and --data that are trained together.'
+        if not isinstance(imgsz, (list, tuple)):
+            model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz))  # warmup
+
+        pad, rect = (0.0, False) if task == 'speed' else (0.5, pt)  # square inference for benchmarks
+
+        # SYNET MODIFICATION: if tflite, use rect with no padding
+        if tflite:
+            pad, rect = 0.0, True
+            stride = np.gcd(sn[2], sn[1])
+
+        task = task if task in ('train', 'val', 'test') else 'val'  # path to train/val/test images
+        dataloader = create_dataloader(data[task],
+                                       imgsz,
+                                       batch_size,
+                                       stride,
+                                       single_cls,
+                                       pad=pad,
+                                       rect=rect,
+                                       workers=workers,
+                                       prefix=colorstr(f'{task}: '))[0]
+
+    seen = 0
+    confusion_matrix = ConfusionMatrix(nc=nc)
+    names = model.names if hasattr(model, 'names') else model.module.names  # get class names
+    if isinstance(names, (list, tuple)):  # old format
+        names = dict(enumerate(names))
+    class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
+    s = ('%22s' + '%11s' * 6) % ('Class', 'Images', 'Instances', 'P', 'R', 'mAP50', 'mAP50-95')
+    tp, fp, p, r, f1, mp, mr, map50, ap50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
+    dt = Profile(), Profile(), Profile()  # profiling times
+    loss = torch.zeros(3, device=device)
+    jdict, stats, ap, ap_class = [], [], [], []
+    callbacks.run('on_val_start')
+    pbar = tqdm(dataloader, desc=s, bar_format=TQDM_BAR_FORMAT)  # progress bar
+    for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
+        callbacks.run('on_val_batch_start')
+        with dt[0]:
+            if cuda:
+                im = im.to(device, non_blocking=True)
+                targets = targets.to(device)
+            im = im.half() if half else im.float()  # uint8 to fp16/32
+            im /= 255  # 0 - 255 to 0.0 - 1.0
+            nb, _, height, width = im.shape  # batch size, channels, height, width
+
+        # SYNET MODIFICATION: if tflite, make grayscale
+        if tflite:
+            im = im.mean(1, keepdims=True)
+
+        # Inference
+        with dt[1]:
+            preds, train_out = model(im) if compute_loss else (model(im, augment=augment), None)
+
+        # Loss
+        if compute_loss:
+            loss += compute_loss(train_out, targets)[1]  # box, obj, cls
+
+        # NMS
+        targets[:, 2:] *= torch.tensor((width, height, width, height), device=device)  # to pixels
+        lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []  # for autolabelling
+        with dt[2]:
+            preds = non_max_suppression(preds,
+                                        conf_thres,
+                                        iou_thres,
+                                        labels=lb,
+                                        multi_label=True,
+                                        agnostic=single_cls,
+                                        max_det=max_det)
+
+        # Metrics
+        for si, pred in enumerate(preds):
+            labels = targets[targets[:, 0] == si, 1:]
+            nl, npr = labels.shape[0], pred.shape[0]  # number of labels, predictions
+            path, shape = Path(paths[si]), shapes[si][0]
+            correct = torch.zeros(npr, niou, dtype=torch.bool, device=device)  # init
+            seen += 1
+
+            if npr == 0:
+                if nl:
+                    stats.append((correct, *torch.zeros((2, 0), device=device), labels[:, 0]))
+                    if plots:
+                        confusion_matrix.process_batch(detections=None, labels=labels[:, 0])
+                continue
+
+            # Predictions
+            if single_cls:
+                pred[:, 5] = 0
+            predn = pred.clone()
+            scale_boxes(im[si].shape[1:], predn[:, :4], shape, shapes[si][1])  # native-space pred
+
+            # Evaluate
+            if nl:
+                tbox = xywh2xyxy(labels[:, 1:5])  # target boxes
+                scale_boxes(im[si].shape[1:], tbox, shape, shapes[si][1])  # native-space labels
+                labelsn = torch.cat((labels[:, 0:1], tbox), 1)  # native-space labels
+                correct = process_batch(predn, labelsn, iouv)
+                if plots:
+                    confusion_matrix.process_batch(predn, labelsn)
+            stats.append((correct, pred[:, 4], pred[:, 5], labels[:, 0]))  # (correct, conf, pcls, tcls)
+
+            # Save/log
+            if save_txt:
+                save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / f'{path.stem}.txt')
+            if save_json:
+                save_one_json(predn, jdict, path, class_map)  # append to COCO-JSON dictionary
+            callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
+
+        # Plot images
+        if plots and batch_i < 3:
+            plot_images(im, targets, paths, save_dir / f'val_batch{batch_i}_labels.jpg', names)  # labels
+            plot_images(im, output_to_target(preds), paths, save_dir / f'val_batch{batch_i}_pred.jpg', names)  # pred
+
+        callbacks.run('on_val_batch_end', batch_i, im, targets, paths, shapes, preds)
+
+    # Compute metrics
+    stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)]  # to numpy
+    if len(stats) and stats[0].any():
+        tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
+        ap50, ap = ap[:, 0], ap.mean(1)  # AP@0.5, AP@0.5:0.95
+        mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
+    nt = np.bincount(stats[3].astype(int), minlength=nc)  # number of targets per class
+
+    # Print results
+    pf = '%22s' + '%11i' * 2 + '%11.3g' * 4  # print format
+    LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
+    if nt.sum() == 0:
+        LOGGER.warning(f'WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels')
+
+    # Print results per class
+    if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
+        for i, c in enumerate(ap_class):
+            LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
+
+    # Export results as html
+    header = "Class Images Labels P R mAP@.5 mAP@.5:.95"
+    headers = header.split()
+    data = []
+    data.append(['all', seen, nt.sum(), f"{float(mp):0.3f}", f"{float(mr):0.3f}", f"{float(map50):0.3f}", f"{float(map):0.3f}"])
+    for i, c in enumerate(ap_class):
+        data.append([names[c], seen, nt[c], f"{float(p[i]):0.3f}", f"{float(r[i]):0.3f}", f"{float(ap50[i]):0.3f}", f"{float(ap[i]):0.3f}"])
+    results_df = pd.DataFrame(data,columns=headers)
+    results_html = results_df.to_html()
+    text_file = open(save_dir / "results.html", "w")
+    text_file.write(results_html)
+    text_file.close()
+
+    # Print speeds
+    t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
+    if not training:
+        if isinstance(imgsz, (list, tuple)):
+            shape = (batch_size, 3, *imgsz)
+        else:
+            shape = (batch_size, 3, imgsz, imgsz)
+        LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)
+
+    # Plots
+    if plots:
+        confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
+        callbacks.run('on_val_end', nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix)
+
+    # Save JSON
+    if save_json and len(jdict):
+        w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ''  # weights
+        anno_json = str(Path('../datasets/coco/annotations/instances_val2017.json'))  # annotations
+        pred_json = str(save_dir / f"{w}_predictions.json")  # predictions
+        LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
+        with open(pred_json, 'w') as f:
+            json.dump(jdict, f)
+
+        try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
+            check_requirements('pycocotools>=2.0.6')
+            from pycocotools.coco import COCO
+            from pycocotools.cocoeval import COCOeval
+
+            anno = COCO(anno_json)  # init annotations api
+            pred = anno.loadRes(pred_json)  # init predictions api
+            eval = COCOeval(anno, pred, 'bbox')
+            if is_coco:
+                eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files]  # image IDs to evaluate
+            eval.evaluate()
+            eval.accumulate()
+            eval.summarize()
+            map, map50 = eval.stats[:2]  # update results (mAP@0.5:0.95, mAP@0.5)
+        except Exception as e:
+            LOGGER.info(f'pycocotools unable to run: {e}')
+
+    # Return results
+    model.float()  # for training
+    if not training:
+        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
+        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
+    maps = np.zeros(nc) + map
+    for i, c in enumerate(ap_class):
+        maps[c] = ap[i]
+    map50s = np.zeros(nc) + map50
+    for i, c in enumerate(ap_class):
+        map50s[c] = ap50[i]
+    return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, map50s, t
+
+
+def patch_yolov5(chip=None):
+    """Apply modifications to YOLOv5 for synet"""
+
+    # enable the  chip if given
+    if chip is not None:
+        module = import_module(f"..{chip}", __name__)
+        setattr(yolo, chip, module)
+        yolo.Concat = module.Cat
+        yolo.Detect = module.Detect = Detect
+
+    # use modified val run function for tflites
+    val.run = val_run_tflite
+
+    # yolo uses uint8.  Change to int8
+    common.np = SimpleNamespace(**vars(numpy))
+    common.np.uint8 = common.np.int8
+
+    import synet
+    synet.get_model_backend = get_yolov5_model

synet/base.py

@@ -0,0 +1,1104 @@
+"""base.py is the "export" layer of synet.  As such, it includes the
+logic of how to run as a keras model.  This is handled by cheking the
+'askeras' context manager, and running in "keras mode" if that context
+is enabled.  As a rule of thumb to differentiate between base.py,
+layers.py:
+
+- base.py should only import from torch, keras, and tensorflow.
+- layers.py should only import from base.py.
+"""
+
+from typing import Tuple, Union, Optional, List
+from torch import cat as torch_cat, minimum, tensor, no_grad, empty
+from torch.nn import (Module as Torch_Module,
+                      Conv2d as Torch_Conv2d,
+                      BatchNorm2d as Torch_Batchnorm,
+                      ModuleList,
+                      ReLU as Torch_ReLU,
+                      ConvTranspose2d as Torch_ConvTranspose2d,
+                      Upsample as Torch_Upsample,
+                      AdaptiveAvgPool2d as Torch_AdaptiveAvgPool,
+                      Dropout as Torch_Dropout,
+                      Linear as Torch_Linear)
+from torch.nn.functional import pad
+import torch.nn as nn
+import torch
+
+
+class AsKeras:
+    """AsKeras is a context manager used to export from pytorch to
+keras.  See test.py and quantize.py for examples.
+
+    """
+
+    def __init__(self):
+        self.use_keras = False
+        self.kwds = dict(train=False)
+
+    def __call__(self, **kwds):
+        self.kwds.update(kwds)
+        return self
+
+    def __enter__(self):
+        self.use_keras = True
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self.__init__()
+
+
+askeras = AsKeras()
+
+
+class Module(Torch_Module):
+    def forward(self, x):
+        if askeras.use_keras and hasattr(self, 'as_keras'):
+            return self.as_keras(x)
+        return self.module(x)
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError as e:
+            if name == 'module':
+                raise e
+            return getattr(self.module, name)
+
+    def to_keras(self, imgsz, in_channels=1, batch_size=1, **kwds):
+        from keras import Input, Model
+        inp = Input(list(imgsz) + [in_channels], batch_size=batch_size)
+        with askeras(imgsz=imgsz, **kwds):
+            return Model(inp, self(inp))
+
+
+class Conv2d(Module):
+    """Convolution operator which ensures padding is done equivalently
+    between PyTorch and TensorFlow.
+
+    """
+
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: Union[int, Tuple[int, int]],
+                 stride: int = 1,
+                 bias: bool = False,
+                 padding: Optional[bool] = True,
+                 groups: Optional[int] = 1):
+        """
+        Implementation of torch Conv2D with option fot supporting keras
+            inference
+        :param in_channels: Number of channels in the input
+        :param out_channels: Number of channels produced by the convolution
+        :param kernel_size: Size of the kernel
+        :param stride:
+        :param bias:
+        :param groups: using for pointwise/depthwise
+        """
+        super().__init__()
+        if isinstance(kernel_size, int):
+            kernel_size = (kernel_size, kernel_size)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = "same" if padding else 'valid'
+        self.groups = groups
+        self.conv = Torch_Conv2d(in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 kernel_size=kernel_size,
+                                 bias=bias,
+                                 stride=stride,
+                                 groups=self.groups)
+        self.use_bias = bias
+
+    def forward(self, x):
+
+        # temporary code for backwards compatibility
+        if not hasattr(self, 'padding'):
+            self.padding = 'same'
+        if not hasattr(self, 'groups'):
+            self.groups = 1
+        if not isinstance(self.padding, str):
+            self.padding = "same" if self.padding else 'valid'
+
+        if askeras.use_keras:
+            return self.as_keras(x)
+
+        if self.padding == "valid":
+            return self.conv(x)
+
+        # make padding like in tensorflow, which right aligns convolutionn.
+        H, W = (s if isinstance(s, int) else s.item() for s in x.shape[-2:])
+        # radius of the kernel and carry.  Border size + carry.  All in y
+        ry, rcy = divmod(self.kernel_size[0] - 1, 2)
+        by, bcy = divmod((H - 1) % self.stride - rcy, 2)
+        # radius of the kernel and carry.  Border size + carry.  All in x
+        rx, rcx = divmod(self.kernel_size[1] - 1, 2)
+        bx, bcx = divmod((W - 1) % self.stride - rcx, 2)
+        # apply pad
+        return self.conv(
+            pad(x, (rx - bx - bcx, rx - bx, ry - by - bcy, ry - by)))
+
+    def as_keras(self, x):
+        if askeras.kwds.get('demosaic'):
+            from .demosaic import Demosaic, reshape_conv
+            demosaic = Demosaic(*askeras.kwds['demosaic'].split('-'))
+            del askeras.kwds['demosaic']
+            return reshape_conv(self)(demosaic(x))
+        from keras.layers import Conv2D as Keras_Conv2d
+        assert x.shape[-1] == self.in_channels, (x.shape, self.in_channels)
+        conv = Keras_Conv2d(filters=self.out_channels,
+                            kernel_size=self.kernel_size,
+                            strides=self.stride,
+                            padding=self.padding,
+                            use_bias=self.use_bias,
+                            groups=self.groups)
+        conv.build(x.shape)
+        if isinstance(self.conv, Torch_Conv2d):
+            tconv = self.conv
+        else:
+            # for NNI compatibility
+            tconv = self.conv.module
+        weight = tconv.weight.detach().numpy().transpose(2, 3, 1, 0)
+        conv.set_weights([weight, tconv.bias.detach().numpy()]
+                         if self.use_bias else
+                         [weight])
+        return conv(x)
+
+    def requires_grad_(self, val):
+        self.conv = self.conv.requires_grad_(val)
+        return self
+
+    def __getattr__(self, name):
+        if name in ("bias", "weight"):
+            return getattr(self.conv, name)
+        return super().__getattr__(name)
+
+    def __setattr__(self, name, value):
+        if name in ("bias", "weight"):
+            return setattr(self.conv, name, value)
+        return super().__setattr__(name, value)
+
+    def split_channels(self, chans):
+
+        with no_grad():
+            split = Conv2d(self.in_channels, len(chans),
+                           self.kernel_size, self.stride, self.use_bias)
+            split.weight[:] = self.weight[chans]
+
+            rest_chans = [i for i in range(self.out_channels)
+                          if i not in chans]
+            rest = Conv2d(self.in_channels, self.out_channels - len(chans),
+                          self.kernel_size, self.stride, self.use_bias)
+            rest.weight[:] = self.weight[rest_chans]
+
+            if self.use_bias:
+                split.bias[:] = self.bias[chans]
+                rest.bias[:] = self.bias[rest_chans]
+
+        return split, rest
+
+
+# don't try to move this assignment into class def.  It won't work.
+# This is for compatibility with NNI so it does not treat this like a
+# pytorch conv2d, and instead finds the nested conv2d.
+Conv2d.__name__ = "Synet_Conv2d"
+
+class DepthwiseConv2d(Conv2d):
+    """DepthwiseConv2d operator implemented as a group convolution for
+       pytorch and DepthwiseConv2d operator for keras
+    """
+    def as_keras(self, x):
+        if askeras.kwds.get('demosaic'):
+            from .demosaic import Demosaic, reshape_conv
+            demosaic = Demosaic(*askeras.kwds['demosaic'].split('-'))
+            del askeras.kwds['demosaic']
+            return reshape_conv(self)(demosaic(x))
+        from keras.layers import DepthwiseConv2D as Keras_DWConv2d
+        assert x.shape[-1] == self.in_channels, (x.shape, self.in_channels)
+        conv = Keras_DWConv2d(kernel_size=self.kernel_size,
+                              strides=self.stride,
+                              padding=self.padding,
+                              use_bias=self.use_bias)
+        conv.build(x.shape)
+        if isinstance(self.conv, Torch_Conv2d):
+            tconv = self.conv
+        else:
+            # for NNI compatibility
+            tconv = self.conv.module
+        weight = tconv.weight.detach().numpy().transpose(2, 3, 0, 1)
+        conv.set_weights([weight, tconv.bias.detach().numpy()]
+                         if self.use_bias else
+                         [weight])
+        return conv(x)
+
+DepthwiseConv2d.__name__ = "Synet_DepthwiseConv2d"
+
+class ConvTranspose2d(Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding,
+                 bias=False):
+        print("WARNING: synet ConvTranspose2d mostly untested")
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = "valid" if padding == 0 else "same"
+        self.use_bias = bias
+        self.module = Torch_ConvTranspose2d(in_channels, out_channels,
+                                            kernel_size, stride,
+                                            padding, bias=bias)
+
+    def as_keras(self, x):
+        from keras.layers import Conv2DTranspose as Keras_ConvTrans
+        conv = Keras_ConvTrans(self.out_channels, self.kernel_size, self.stride,
+                               self.padding, use_bias=self.use_bias)
+        conv.build(x.shape)
+        if isinstance(self.module, Torch_ConvTranspose2d):
+            tconv = self.module
+        else:
+            # for NNI compatibility
+            tconv = self.module.module
+        weight = tconv.weight.detach().numpy().transpose(2, 3, 1, 0)
+        conv.set_weights([weight, tconv.bias.detach().numpy()]
+                         if self.use_bias else
+                         [weight])
+        return conv(x)
+
+
+class Cat(Module):
+    """Concatenate along feature dimension."""
+
+    def __init__(self, *args):
+        super().__init__()
+
+    def forward(self, xs):
+        if askeras.use_keras:
+            return self.as_keras(xs)
+        return torch_cat(xs, dim=1)
+
+    def as_keras(self, xs):
+        assert all(len(x.shape) == 4 for x in xs)
+        from keras.layers import Concatenate as Keras_Concatenate
+        return Keras_Concatenate(-1)(xs)
+
+
+class ReLU(Module):
+    def __init__(self, max_val=None, name=None):
+        super().__init__()
+        self.max_val = None if max_val is None else tensor(max_val,
+                                                           dtype=float)
+        self.name = name
+        self.relu = Torch_ReLU()
+
+    def forward(self, x):
+        if askeras.use_keras:
+            return self.as_keras(x)
+        if self.max_val is None:
+            return self.relu(x)
+        return minimum(self.relu(x), self.max_val)
+
+    def as_keras(self, x):
+        # temporary code for backwards compatibility
+        if not hasattr(self, 'name'):
+            self.name = None
+        from keras.layers import ReLU as Keras_ReLU
+        return Keras_ReLU(self.max_val, name=self.name)(x)
+
+
+class BatchNorm(Module):
+    def __init__(self, features, epsilon=1e-3, momentum=0.999):
+        super().__init__()
+        self.epsilon = epsilon
+        self.momentum = momentum
+        self.module = Torch_Batchnorm(features, epsilon, momentum)
+
+    def forward(self, x):
+        # temporary code for backwards compatibility
+        if hasattr(self, 'batchnorm'):
+            self.module = self.batchnorm
+        return super().forward(x)
+
+    def as_keras(self, x):
+
+        from keras.layers import BatchNormalization as Keras_Batchnorm
+        batchnorm = Keras_Batchnorm(momentum=self.momentum,
+                                    epsilon=self.epsilon)
+        batchnorm.build(x.shape)
+        if isinstance(self.module, Torch_Batchnorm):
+            bn = self.module
+        else:
+            bn = self.module.module
+        weights = bn.weight.detach().numpy()
+        bias = bn.bias.detach().numpy()
+        running_mean = bn.running_mean.detach().numpy()
+        running_var = bn.running_var.detach().numpy()
+        batchnorm.set_weights([weights, bias, running_mean, running_var])
+        return batchnorm(x, training=askeras.kwds["train"])
+
+
+class Upsample(Module):
+    allowed_modes = "bilinear", "nearest"
+
+    def __init__(self, scale_factor, mode="nearest"):
+        assert mode in self.allowed_modes
+        if not isinstance(scale_factor, int):
+            for sf in scale_factor:
+                assert isinstance(sf, int)
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.mode = mode
+        self.module = Torch_Upsample(scale_factor=scale_factor, mode=mode)
+
+    def forward(self, x):
+        # temporary code for backwards compatibility
+        if not hasattr(self, 'module'):
+            self.module = self.upsample
+        return super().forward(x)
+
+    def as_keras(self, x):
+        from keras.layers import UpSampling2D
+        return UpSampling2D(size=self.scale_factor,
+                            interpolation=self.mode,
+                            )(x)
+
+
+class Sequential(Module):
+    def __init__(self, *sequence):
+        super().__init__()
+        self.ml = ModuleList(sequence)
+
+    def forward(self, x):
+        for layer in self.ml:
+            x = layer(x)
+        return x
+
+    def __getitem__(self, i):
+        if isinstance(i, int):
+            return self.ml[i]
+        return Sequential(*self.ml[i])
+
+
+class GlobalAvgPool(Module):
+    def __init__(self):
+        super().__init__()
+        self.module = Torch_AdaptiveAvgPool(1)
+
+    def as_keras(self, x):
+        from keras.layers import GlobalAveragePooling2D
+        return GlobalAveragePooling2D(keepdims=True)(x)
+
+
+class Dropout(Module):
+    def __init__(self, p=0, inplace=False):
+        super().__init__()
+        self.p = p
+        self.module = Torch_Dropout(p, inplace=inplace)
+
+    def as_keras(self, x):
+        from keras.layers import Dropout
+        return Dropout(self.p)(x, training=askeras.kwds["train"])
+
+
+class Linear(Module):
+    def __init__(self, in_c, out_c, bias=True):
+        super().__init__()
+        self.use_bias = bias
+        self.module = Torch_Linear(in_c, out_c, bias)
+
+    def as_keras(self, x):
+        from keras.layers import Dense
+        out_c, in_c = self.module.weight.shape
+        params = [self.module.weight.detach().numpy().transpose(1, 0)]
+        if self.use_bias:
+            params.append(self.module.bias.detach().numpy())
+        dense = Dense(out_c, use_bias=self.use_bias)
+        dense.build(x.shape[1:])
+        dense.set_weights(params)
+        return dense(x)
+
+
+class Transpose(Module):
+    """
+    A class designed to transpose tensors according to specified dimension permutations, compatible
+    with both PyTorch and TensorFlow (Keras). It allows for flexible tensor manipulation, enabling
+    dimension reordering to accommodate the requirements of different neural network architectures
+    or operations.
+
+    The class supports optional channel retention during transposition in TensorFlow to ensure
+    compatibility with Keras' channel ordering conventions.
+    """
+
+    def forward(self, x, perm: Union[Tuple[int], List[int]],
+                keep_channel_last: bool = False):
+        """
+        Transposes the input tensor according to the specified dimension permutation. If integrated
+        with Keras, it converts PyTorch tensors to TensorFlow tensors before transposing, with an
+        option to retain channel ordering as per Keras convention.
+
+        Parameters:
+            x (Tensor): The input tensor to be transposed.
+            perm (tuple or list): The permutation of dimensions to apply to the tensor.
+            keep_channel_last (bool, optional): Specifies whether to adjust the permutation to
+                                                 retain Keras' channel ordering convention. Default
+                                                 is False.
+
+        Returns:
+            Tensor: The transposed tensor.
+        """
+        if askeras.use_keras:
+            return self.as_keras(x, perm, keep_channel_last)
+        # Use PyTorch's permute method for the operation
+        return x.permute(*perm)
+
+    def as_keras(self, x, perm: Union[Tuple[int], List[int]],
+                 keep_channel_last: bool):
+        """
+        Handles tensor transposition in a TensorFlow/Keras environment, converting PyTorch tensors
+        to TensorFlow tensors if necessary, and applying the specified permutation. Supports an
+        option for channel retention according to Keras conventions.
+
+        Parameters:
+            x (Tensor): The input tensor, possibly a PyTorch tensor.
+            perm (tuple or list): The permutation of dimensions to apply.
+            keep_channel_last (bool): If True, adjusts the permutation to retain Keras' channel
+                                       ordering convention.
+
+        Returns:
+            Tensor: The transposed tensor in TensorFlow format.
+        """
+        import tensorflow as tf
+
+        # Adjust for TensorFlow's default channel ordering if necessary
+        tf_format = [0, 3, 1, 2]
+
+        # Map PyTorch indices to TensorFlow indices if channel retention is enabled
+        mapped_indices = [tf_format[index] for index in
+                          perm] if keep_channel_last else perm
+
+        # Convert PyTorch tensors to TensorFlow tensors if necessary
+        x_tf = tf.convert_to_tensor(x.detach().numpy(),
+                                    dtype=tf.float32) if isinstance(x,
+                                                                    torch.Tensor) else x
+
+        # Apply the transposition with TensorFlow's transpose method
+        x_tf_transposed = tf.transpose(x_tf, perm=mapped_indices)
+
+        return x_tf_transposed
+
+
+class Reshape(Module):
+    """
+    A class designed to reshape tensors to a specified shape, compatible with both PyTorch and
+    TensorFlow (Keras). This class facilitates tensor manipulation across different deep learning
+    frameworks, enabling the adjustment of tensor dimensions to meet the requirements of different
+    neural network layers or operations.
+
+    It supports dynamic reshaping capabilities, automatically handling the conversion between
+    PyTorch and TensorFlow tensors and applying the appropriate reshaping operation based on the
+    runtime context.
+    """
+
+    def forward(self, x, shape: Union[Tuple[int], List[int]]):
+        """
+        Reshapes the input tensor to the specified shape. If integrated with Keras, it converts
+        PyTorch tensors to TensorFlow tensors before reshaping.
+
+        Parameters:
+            x (Tensor): The input tensor to be reshaped.
+            shape (tuple or list): The new shape for the tensor. The specified shape can include
+                                   a `-1` to automatically infer the dimension that ensures the
+                                   total size remains constant.
+
+        Returns:
+            Tensor: The reshaped tensor.
+        """
+        if askeras.use_keras:
+            return self.as_keras(x, shape)
+        # Use PyTorch's reshape method for the operation
+        return x.reshape(*shape)
+
+    def as_keras(self, x, shape: Union[Tuple[int], List[int]]):
+        """
+        Converts PyTorch tensors to TensorFlow tensors, if necessary, and performs the reshape
+        operation using TensorFlow's reshape function. This method ensures compatibility and
+        functionality within a TensorFlow/Keras environment.
+
+        Parameters:
+            x (Tensor): The input tensor, possibly a PyTorch tensor.
+            shape (tuple or list): The new shape for the tensor, including the possibility
+                                   of using `-1` to infer a dimension automatically.
+
+        Returns:
+            Tensor: The reshaped tensor in TensorFlow format.
+        """
+        import tensorflow as tf
+        # Convert PyTorch tensors to TensorFlow tensors if necessary
+        x_tf = tf.convert_to_tensor(x.detach().numpy(),
+                                    dtype=tf.float32) if isinstance(x,
+                                                                    torch.Tensor) else x
+
+        # Use TensorFlow's reshape function to adjust the tensor's dimensions
+        x_tf_reshaped = tf.reshape(x_tf, shape)
+
+        # TensorFlow's reshape might introduce an additional dimension if shape is fully defined,
+        # use tf.squeeze to adjust dimensions if necessary
+        return x_tf_reshaped
+
+
+class Flip(Module):
+    """
+    A class to flip tensors along specified dimensions, supporting both PyTorch and TensorFlow
+    (Keras). This class enables consistent tensor manipulation across different deep learning
+    frameworks, facilitating operations like data augmentation or image processing where flipping
+    is required.
+
+    The class automatically detects the runtime environment to apply the appropriate flipping
+    operation, handling tensor conversions between PyTorch and TensorFlow as needed.
+    """
+
+    def forward(self, x, dims: Union[List[int], Tuple[int]]):
+        """
+        Flips the input tensor along specified dimensions. If integrated with Keras, it
+        converts PyTorch tensors to TensorFlow tensors before flipping.
+
+        Parameters:
+            x (Tensor): The input tensor to be flipped.
+            dims (list or tuple): The dimensions along which to flip the tensor.
+
+        Returns:
+            Tensor: The flipped tensor.
+        """
+        # Check if Keras usage is flagged and handle accordingly
+        if askeras.use_keras:
+            return self.as_keras(x, dims)
+        # Use PyTorch's flip function for the operation
+        return torch.flip(x, dims)
+
+    def as_keras(self, x, dims: Union[List[int], Tuple[int]]):
+        """
+        Converts PyTorch tensors to TensorFlow tensors, if necessary, and performs the flip
+        operation using TensorFlow's reverse function. This method ensures compatibility and
+        functionality within a TensorFlow/Keras environment.
+
+        Parameters:
+            x (Tensor): The input tensor, possibly a PyTorch tensor.
+            dims (list or tuple): The dimensions along which to flip the tensor.
+
+        Returns:
+            Tensor: The flipped tensor in TensorFlow format.
+        """
+        import tensorflow as tf
+        # Convert PyTorch tensors to TensorFlow tensors if necessary
+        x_tf = tf.convert_to_tensor(x.detach().numpy(),
+                                    dtype=tf.float32) if isinstance(x,
+                                                                    torch.Tensor) else x
+
+        # Use TensorFlow's reverse function for flipping along specified dimensions
+        return tf.reverse(x_tf, axis=dims)
+
+
+class Add(Module):
+    """
+    A class designed to perform element-wise addition on tensors, compatible with both
+    PyTorch and TensorFlow (Keras). This enables seamless operation across different deep
+    learning frameworks, supporting the addition of tensors regardless of their originating
+    framework.
+
+    The class automatically handles framework-specific tensor conversions and uses the
+    appropriate addition operation based on the runtime context, determined by whether
+    TensorFlow/Keras or PyTorch is being used.
+    """
+
+    def forward(self, x, y):
+        """
+        Performs element-wise addition of two tensors. If integrated with Keras, converts
+        PyTorch tensors to TensorFlow tensors before addition.
+
+        Parameters:
+            x (Tensor): The first input tensor.
+            y (Tensor): The second input tensor to be added to the first.
+
+        Returns:
+            Tensor: The result of element-wise addition of `x` and `y`.
+        """
+        if askeras.use_keras:
+            return self.as_keras(x, y)
+        # Use PyTorch's add function for element-wise addition
+        return torch.add(x, y)
+
+    def as_keras(self, x, y):
+        """
+        Converts PyTorch tensors to TensorFlow tensors, if necessary, and performs
+        element-wise addition using TensorFlow's add function. This method ensures
+        compatibility and functionality within a TensorFlow/Keras environment.
+
+        Parameters:
+            x (Tensor): The first input tensor, possibly a PyTorch tensor.
+            y (Tensor): The second input tensor, possibly a PyTorch tensor.
+
+        Returns:
+            Tensor: The result of element-wise addition of `x` and `y` in TensorFlow format.
+        """
+        import tensorflow as tf
+        # Convert PyTorch tensors to TensorFlow tensors if necessary
+        x_tf = tf.convert_to_tensor(x.detach().numpy(),
+                                    dtype=tf.float32) if isinstance(x,
+                                                                    torch.Tensor) else x
+        y_tf = tf.convert_to_tensor(y.detach().numpy(),
+                                    dtype=tf.float32) if isinstance(y,
+                                                                    torch.Tensor) else y
+
+        # Use TensorFlow's add function for element-wise addition
+        return tf.add(x_tf, y_tf)
+
+
+class Shape(Module):
+    """
+    A utility class for obtaining the shape of a tensor in a format compatible
+    with either PyTorch or Keras. This class facilitates the transformation of
+    tensor shapes, particularly useful for adapting model input or output
+    dimensions across different deep learning frameworks.
+
+    The class provides a method to directly return the shape of a tensor for
+    PyTorch use cases and an additional method for transforming the shape to a
+    Keras-compatible format, focusing on the common difference in dimension
+    ordering between the two frameworks.
+    """
+
+    def forward(self, x):
+        """
+        Returns the shape of the tensor. If integrated with Keras, it transforms the tensor shape
+        to be compatible with Keras dimension ordering.
+
+        Parameters:
+            x (Tensor): The input tensor whose shape is to be obtained or transformed.
+
+        Returns:
+            Tuple: The shape of the tensor, directly returned for PyTorch or transformed for Keras.
+        """
+        if askeras.use_keras:
+            return self.as_keras(x)
+        # Directly return the shape for PyTorch tensors
+        return x.shape
+
+    def as_keras(self, x):
+        """
+        Transforms the tensor shape to be compatible with Keras' expected dimension ordering.
+        This method is designed to switch between CHW and HWC formats based on the tensor's
+        dimensionality, handling common cases for 2D, 3D, and 4D tensors.
+
+        Parameters:
+            x (Tensor): The input tensor whose shape is to be transformed for Keras.
+
+        Returns:
+            Tuple: The transformed shape of the tensor, suitable for Keras models.
+        """
+        # Handle different tensor dimensionality with appropriate
+        # transformations
+        if len(x.shape) == 4:  # Assuming NCHW format, convert to NHWC
+            N, W, H, C = x.shape
+            x_shape = (N, C, H, W)
+        elif len(x.shape) == 3:  # Assuming CHW format, convert to HWC
+            H, W, C = x.shape
+            x_shape = (C, H, W)
+        else:  # Assuming 2D tensor, no channel dimension involved
+            H, W = x.shape
+            x_shape = (H, W)
+
+        return x_shape
+
+
+class GenericRNN(nn.Module):
+    """
+    A base class for customizable RNN models supporting RNN, GRU, and LSTM networks.
+    """
+
+    def __init__(self, input_size: int, hidden_size: int, num_layers: int = 1,
+                 bidirectional: bool = False, bias: bool = True,
+                 batch_first: bool = True, dropout: float = 0) -> None:
+        super(GenericRNN, self).__init__()
+        self.bidirectional = 2 if bidirectional else 1
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bias = bias
+        self.dropout = dropout
+        self.input_size = input_size
+        self.batch_first = batch_first
+
+    def init_rnn(self, input_size: int, hidden_size: int, num_layers: int,
+                 bidirectional: bool, bias: bool, batch_first: bool,
+                 dropout: float) -> None:
+
+        raise NotImplementedError("Must be implemented by subclass.")
+
+    def forward(self, x, h0=None, c0=None):
+
+        raise NotImplementedError("Must be implemented by subclass.")
+
+    def as_keras(self, x):
+        raise NotImplementedError("Must be implemented by subclass.")
+
+    def generic_as_keras(self, x, RNNBase):
+        """
+        Converts the model architecture and weights to a Keras-compatible format and applies
+        the model to the provided input.
+
+        This method enables the use of PyTorch-trained models within the Keras framework by
+        converting the input tensor to a TensorFlow tensor, recreating the model architecture
+        in Keras, and setting the weights accordingly.
+
+        Parameters:
+            x (Tensor): The input tensor for the model, which can be a PyTorch tensor.
+            RNNBase: The base class for the RNN model to be used in the Keras model.
+
+        Returns:
+            Tuple[Tensor, None]: A tuple containing the output of the Keras model applied to the
+                                 converted input tensor and None (since Keras models do not
+                                 necessarily return the final hidden state as PyTorch models do).
+
+        Raises:
+            ImportError: If the required TensorFlow or Keras modules are not available.
+        """
+
+        # Import necessary modules from Keras and TensorFlow
+        from keras.layers import Bidirectional
+        from keras.models import Sequential as KerasSequential
+        import tensorflow as tf
+
+        # Convert PyTorch tensor to TensorFlow tensor if necessary
+        if isinstance(x, torch.Tensor):
+            x_tf = tf.convert_to_tensor(x.detach().numpy(), dtype=tf.float32)
+        else:
+            x_tf = x
+
+        # Create a Keras Sequential model for stacking layers
+        model = KerasSequential()
+
+        # Add RNN layers to the Keras model
+        for i in range(self.num_layers):
+            # Determine if input shape needs to be specified (only for the first layer)
+            if i == 0:
+                layer = RNNBase(units=self.hidden_size, return_sequences=True,
+                                input_shape=list(x.shape[1:]),
+                                use_bias=self.bias,
+                                dropout=self.dropout if i < self.num_layers - 1 else 0)
+            else:
+                layer = RNNBase(units=self.hidden_size, return_sequences=True,
+                                use_bias=self.bias,
+                                dropout=self.dropout if i < self.num_layers - 1 else 0)
+
+            # Wrap the layer with Bidirectional if needed
+            if self.bidirectional == 2:
+                layer = Bidirectional(layer)
+
+            model.add(layer)
+
+        # Apply previously extracted PyTorch weights to the Keras model
+        self.set_keras_weights(model)
+
+        # Process the input through the Keras model
+        output = model(x_tf)
+
+        # Return the output and None for compatibility with PyTorch output format
+        return output, None
+
+    def extract_pytorch_rnn_weights(self):
+        """
+        Extracts weights from a PyTorch model's RNN layers and prepares them for
+        transfer to a Keras model.
+
+        This function iterates through the named parameters of a PyTorch model,
+        detaching them from the GPU (if applicable),
+        moving them to CPU memory, and converting them to NumPy arrays.
+        It organizes these weights in a dictionary,
+        using the parameter names as keys, which facilitates their later use in
+        setting weights for a Keras model.
+
+        Returns:
+        A dictionary containing the weights of the PyTorch model, with parameter
+        names as keys and their corresponding NumPy array representations as values.
+        """
+
+        weights = {}  # Initialize a dictionary to store weights
+
+        # Iterate through the model's named parameters
+        for name, param in self.named_parameters():
+            # Process the parameter name to extract the relevant part
+            # and use it as the key in the weights dictionary
+            key = name.split('.')[
+                -1]  # Extract the last part of the parameter name
+
+            # Detach the parameter from the computation graph, move it to CPU,
+            # and convert to NumPy array
+            weights[key] = param.detach().cpu().numpy()
+
+        return weights  # Return the dictionary of weights
+
+    def set_keras_weights(self, keras_model):
+        raise NotImplementedError("Must be implemented by subclass.")
+
+    def generic_set_keras_weights(self, keras_model, RNNBase: str):
+        """
+        Sets the weights of a Keras model based on the weights from a PyTorch model.
+
+        This function is designed to transfer weights from PyTorch RNN layers (SimpleRNN, GRU, LSTM)
+        to their Keras counterparts, including handling for bidirectional layers. It ensures that the
+        weights are correctly transposed and combined to match Keras's expectations.
+
+        Parameters:
+        - keras_model: The Keras model to update the weights for.
+        """
+
+        # Import necessary modules
+        from keras.layers import Bidirectional
+        import numpy as np
+
+        # Extract weights from PyTorch model
+        pytorch_weights = self.extract_pytorch_rnn_weights()
+
+        # Iterate over each layer in the Keras model
+        for layer in keras_model.layers:
+            # Check if layer is bidirectional and set layers to update
+            # accordingly
+            if isinstance(layer, Bidirectional):
+                layers_to_update = [layer.layer, layer.backward_layer]
+            else:
+                layers_to_update = [layer]
+
+            # Update weights for each RNN layer in layers_to_update
+            for rnn_layer in layers_to_update:
+
+                num_gates = {'SimpleRNN': 1, 'GRU': 3, 'LSTM': 4}.get(
+                    RNNBase, 0)
+
+                # Initialize lists for input-hidden, hidden-hidden weights,
+                # and biases
+                ih_weights, hh_weights, biases = [], [], []
+
+                # Process weights and biases for each gate
+                for i in range(num_gates):
+                    gate_suffix = f'_l{i}'
+                    for prefix in ('weight_ih', 'weight_hh'):
+                        key = f'{prefix}{gate_suffix}'
+                        if key in pytorch_weights:
+                            weights = \
+                                pytorch_weights[key].T  # Transpose to match Keras shape
+
+                            (ih_weights if prefix == 'weight_ih' else hh_weights) \
+                                .append(weights)
+
+                    bias_keys = (
+                        f'bias_ih{gate_suffix}', f'bias_hh{gate_suffix}')
+                    if all(key in pytorch_weights for key in bias_keys):
+                        # Sum biases from input-hidden and hidden-hidden
+                        biases.append(
+                            sum(pytorch_weights[key] for key in bias_keys))
+
+                # Combine weights and biases into a format suitable for Keras
+                keras_weights = [np.vstack(ih_weights),
+                                 np.vstack(hh_weights), np.hstack(biases)]
+
+                # Set the weights for the Keras layer
+                if not isinstance(layer, Bidirectional):
+                    rnn_layer.set_weights(keras_weights)
+                else:
+                    rnn_layer.cell.set_weights(keras_weights)
+
+
+class RNN(GenericRNN):
+    def __init__(self, *args, **kwargs):
+        super(RNN, self).__init__(*args, **kwargs)
+        self.rnn = nn.RNN(input_size=kwargs['input_size'],
+                          hidden_size=kwargs['hidden_size'],
+                          num_layers=kwargs['num_layers'],
+                          bias=kwargs['bias'],
+                          batch_first=kwargs['batch_first'],
+                          dropout=kwargs['dropout'],
+                          bidirectional=kwargs['bidirectional'])
+
+    def forward(self, x, h0=None):
+        if askeras.use_keras:
+            return self.as_keras(x)
+
+        out, h = self.rnn(x, h0)
+        return out, h
+
+    def as_keras(self, x):
+        """
+        Converts the model architecture and weights to a Keras-compatible format and applies
+        the model to the provided input.
+
+        This method enables the use of PyTorch-trained models within the Keras framework by
+        converting the input tensor to a TensorFlow tensor, recreating the model architecture
+        in Keras, and setting the weights accordingly.
+
+        Parameters:
+            x (Tensor): The input tensor for the model, which can be a PyTorch tensor.
+
+        Returns:
+            Tuple[Tensor, None]: A tuple containing the output of the Keras model applied to the
+                                 converted input tensor and None (since Keras models do not
+                                 necessarily return the final hidden state as PyTorch models do).
+
+        Raises:
+            ImportError: If the required TensorFlow or Keras modules are not available.
+        """
+
+        # Import necessary modules from Keras and TensorFlow
+        from keras.layers import SimpleRNN
+
+        output, _ = super().generic_as_keras(x, SimpleRNN)
+
+        return output, None
+
+    def set_keras_weights(self, keras_model):
+        """
+        Sets the weights of a Keras model based on the weights from a PyTorch model.
+
+        This function is designed to transfer weights from PyTorch RNN layers
+        to their Keras counterparts, including handling for bidirectional layers. It ensures that the
+        weights are correctly transposed and combined to match Keras's expectations.
+
+        Parameters:
+        - keras_model: The Keras model to update the weights for.
+        """
+
+        self.generic_set_keras_weights(keras_model, 'SimpleRNN')
+
+
+class GRU(GenericRNN):
+    def __init__(self, *args, **kwargs):
+        super(GRU, self).__init__(*args, **kwargs)
+        self.rnn = nn.GRU(input_size=kwargs['input_size'],
+                          hidden_size=kwargs['hidden_size'],
+                          num_layers=kwargs['num_layers'],
+                          bias=kwargs['bias'],
+                          batch_first=kwargs['batch_first'],
+                          dropout=kwargs['dropout'],
+                          bidirectional=kwargs['bidirectional'])
+
+    def forward(self, x, h0=None):
+        if askeras.use_keras:
+            return self.as_keras(x)
+
+        out, h = self.rnn(x, h0)
+        return out, h
+
+    def as_keras(self, x):
+        """
+        Converts the model architecture and weights to a Keras-compatible format and applies
+        the model to the provided input.
+
+        This method enables the use of PyTorch-trained models within the Keras framework by
+        converting the input tensor to a TensorFlow tensor, recreating the model architecture
+        in Keras, and setting the weights accordingly.
+
+        Parameters:
+            x (Tensor): The input tensor for the model, which can be a PyTorch tensor.
+
+        Returns:
+            Tuple[Tensor, None]: A tuple containing the output of the Keras model applied to the
+                                 converted input tensor and None (since Keras models do not
+                                 necessarily return the final hidden state as PyTorch models do).
+
+        Raises:
+            ImportError: If the required TensorFlow or Keras modules are not available.
+        """
+
+        # Import necessary modules from Keras and TensorFlow
+        from keras.layers import GRU
+
+        output, _ = super().generic_as_keras(x, GRU)
+
+        return output, None
+
+    def set_keras_weights(self, keras_model):
+        """
+        Sets the weights of a Keras model based on the weights from a PyTorch model.
+
+        This function is designed to transfer weights from PyTorch RNN layers
+        to their Keras counterparts, including handling for bidirectional layers. It ensures that the
+        weights are correctly transposed and combined to match Keras's expectations.
+
+        Parameters:
+        - keras_model: The Keras model to update the weights for.
+        """
+
+        self.generic_set_keras_weights(keras_model, 'GRU')
+
+
+class LSTM(GenericRNN):
+    def __init__(self, *args, **kwargs):
+        super(LSTM, self).__init__(*args, **kwargs)
+        self.rnn = nn.GRU(input_size=kwargs['input_size'],
+                          hidden_size=kwargs['hidden_size'],
+                          num_layers=kwargs['num_layers'],
+                          bias=kwargs['bias'],
+                          batch_first=kwargs['batch_first'],
+                          dropout=kwargs['dropout'],
+                          bidirectional=kwargs['bidirectional'])
+
+    def forward(self, x, h0=None, c0=None):
+        if askeras.use_keras:
+            return self.as_keras(x)
+
+        out, h = self.rnn(x, (h0, c0))
+
+        return out, h
+
+    def as_keras(self, x):
+        """
+        Converts the model architecture and weights to a Keras-compatible format and applies
+        the model to the provided input.
+
+        This method enables the use of PyTorch-trained models within the Keras framework by
+        converting the input tensor to a TensorFlow tensor, recreating the model architecture
+        in Keras, and setting the weights accordingly.
+
+        Parameters:
+            x (Tensor): The input tensor for the model, which can be a PyTorch tensor.
+
+        Returns:
+            Tuple[Tensor, None]: A tuple containing the output of the Keras model applied to the
+                                 converted input tensor and None (since Keras models do not
+                                 necessarily return the final hidden state as PyTorch models do).
+
+        Raises:
+            ImportError: If the required TensorFlow or Keras modules are not available.
+        """
+
+        # Import necessary modules from Keras and TensorFlow
+        from keras.layers import LSTM
+
+        output, _ = super().generic_as_keras(x, LSTM)
+
+        return output, None
+
+    def set_keras_weights(self, keras_model):
+        """
+        Sets the weights of a Keras model based on the weights from a PyTorch model.
+
+        This function is designed to transfer weights from PyTorch RNN layers
+        to their Keras counterparts, including handling for bidirectional layers. It ensures that the
+        weights are correctly transposed and combined to match Keras's expectations.
+
+        Parameters:
+        - keras_model: The Keras model to update the weights for.
+        """
+
+        self.generic_set_keras_weights(keras_model, 'LSTM')
+
+
+class ChannelSlice(Module):
+    def __init__(self, slice):
+        super().__init__()
+        self.slice = slice
+
+    def forward(self, x):
+        if askeras.use_keras:
+            return self.as_keras(x)
+        return x[:, self.slice]
+
+    def as_keras(self, x):
+        return x[(len(x.shape)-1)*(slice(None),)+(self.slice,)]

synet/data_subset.py

@@ -0,0 +1,54 @@
+from os.path import splitext
+def get_label(im):
+    return splitext(get_labels(im))[0] + '.txt'
+
+def get_labels(ims):
+    return "/labels/".join(ims.rsplit("/images/", 1))
+
+from argparse import ArgumentParser
+def parse_opt():
+    parser = ArgumentParser()
+    parser.add_argument("--max-bg-ratio", type=float, default=.999)
+    parser.add_argument('old_yaml')
+    parser.add_argument('new_yaml')
+    return parser.parse_args()
+
+
+from os import listdir, makedirs, symlink
+from os.path import join, abspath, isfile
+from random import shuffle
+from yaml import safe_load as load
+def run(old_yaml, new_yaml, max_bg_ratio):
+    old = load(open(old_yaml))
+    new = load(open(new_yaml))
+    l_n = {l:n for n, l in new['names'].items()}
+    old_cls_new = {str(o): str(l_n[l]) for o, l in old['names'].items()
+                   if l in l_n}
+    splits = ['val', 'train']
+    if 'test' in new:
+        splits.append('test')
+    for split in splits:
+        fg = 0
+        background = []
+        for d in new, old:
+            d[split] = join(d.get('path', ''), d[split])
+        makedirs(new[split])
+        makedirs(get_labels(new[split]))
+        for imf in listdir(old[split]):
+            oldim = join(old[split], imf)
+            newim = join(new[split], imf)
+            labels = [" ".join([old_cls_new[parts[0]], parts[1]])
+                      for label in open(oldlb).readlines()
+                      if (parts := label.split(" ", 1))[0] in old_cls_new
+                      ] if isfile(oldlb := get_label(oldim)) else []
+            if not labels:
+                background.append((oldim, newim))
+            else:
+                fg += 1
+                symlink(abspath(oldim), newim)
+                open(get_label(newim), 'w').writelines(labels)
+
+        shuffle(background)
+        background = background[:int(max_bg_ratio * fg / (1 - max_bg_ratio))]
+        for oldim, newim in background:
+            symlink(abspath(oldim), newim)

synet/demosaic.py

@@ -0,0 +1,290 @@
+from torch import zeros, tensor, arange, where, float32, empty
+from numpy import empty as npempty
+
+from .base import Module, Conv2d, askeras
+
+
+class Mosaic(Module):
+    def __init__(self, bayer_pattern, real_keras=False):
+        super().__init__()
+        self.bayer_pattern = tensor(['rgb'.index(c)
+                                     for c in bayer_pattern.lower()])
+        self.rows = tensor([0, 0, 1, 1])
+        self.cols = tensor([0, 1, 0, 1])
+        self.real_keras = real_keras
+
+    def forward(self, x):
+        if askeras.use_keras:
+            if self.real_keras:
+                return self.as_keras(x)
+            return x
+        *b, c, h, w = x.shape
+        y = empty((*b, 1, h, w), dtype=x.dtype, device=x.device)
+        for yoff, xoff, chan in zip(self.rows, self.cols, self.bayer_pattern):
+            y[..., 0, yoff::2, xoff::2] = x[..., chan, yoff::2, xoff::2]
+        return y
+
+    def clf(self, x):
+        y = npempty((*x.shape[:-1], 1), dtype=x.dtype)
+        for yoff, xoff, chan in zip(self.rows, self.cols, self.bayer_pattern):
+            y[..., yoff::2, xoff::2, 0] = x[..., yoff::2, xoff::2, chan]
+        return y
+
+    def as_keras(self, x):
+        B, H, W, C = x.shape
+        from keras.layers import Concatenate, Reshape
+        a, b, c, d = [x[..., int(yoff)::2, int(xoff)::2, int(chan):int(chan)+1]
+                      for yoff, xoff, chan in
+                      zip(self.rows, self.cols, self.bayer_pattern)]
+        return Reshape((H, W, 1))(
+            Concatenate(-2)((
+                Concatenate(-1)((a, b)),
+                Concatenate(-1)((c, d)))))
+
+
+class MosaicGamma(Mosaic):
+
+    def __init__(self, *args, normalized=True, gammas=[], **kwds):
+        super().__init__(*args, **kwds)
+        self.gammas = gammas
+        if normalized:
+            self.gamma_func = self.normalized_gamma
+        else:
+            self.gamma_func = self.unnormalized_gamma
+
+    def normalized_gamma(self, x, gamma):
+        return x**gamma
+
+    def unnormalized_gamma(self, x, gamma):
+        return ((x / 255)**gamma) * 255
+
+    def as_keras(self, x):
+        from keras.layers import Concatenate
+        a, b, c, d = [self.gamma_func(x[..., int(yoff)::2, int(xoff)::2,
+                                        int(chan):int(chan) + 1],
+                                      self.gammas[chan])
+                      for yoff, xoff, chan in
+                      zip(self.rows, self.cols, self.bayer_pattern)]
+        return Concatenate(-2)((
+            Concatenate(-1)((a, b)),
+            Concatenate(-1)((c, d))))
+
+
+class UnfoldedMosaicGamma(MosaicGamma):
+    def as_keras(self, x):
+        B, H, W, C = x.shape
+        from keras.layers import Reshape
+        return Reshape((H, W, 1))(super().as_keras(x))
+
+
+class Demosaic(Module):
+
+    def __init__(self, dfilter, bayer_pattern, *scales):
+        super().__init__()
+        assert bayer_pattern.lower() in ("rggb", "bggr", "grbg", "gbrg")
+        bayer_pattern = tensor(['rgb'.index(c) for c in bayer_pattern.lower()])
+        rows = tensor([0, 0, 1, 1])
+        cols = tensor([0, 1, 0, 1])
+        # assign kernels from specific filter method
+        getattr(self, dfilter+'_init')()
+        # The basic idea is to apply kxk kernels to two consecutive
+        # rows/columns similtaneously, so we will need a (k+1)x(k+1)
+        # kernel to give it the proper receptive field.  For a given
+        # row or column in the 2x2 bayer grid, we need to either slice
+        # the kxk kernel into the first or last k rows/columns of the
+        # (k+1)x(k+1) generated kernel.
+        kslice = slice(None, -1), slice(1, None)
+        weight = zeros(4, 3, self.k+1, self.k+1, requires_grad=False)
+
+        # Set values for which the bayer image IS ground truth.
+        # +self.k//2 because 2x2 bayer is centered in the (k+1)x(k+1)
+        # kernel.
+        weight[arange(4), bayer_pattern, rows+self.k//2, cols+self.k//2] = 1
+
+        # Finishing off red bayer locations
+        r = bayer_pattern == 'rgb'.index('r')
+        slicey, slicex = kslice[rows[r]], kslice[cols[r]]
+        weight[r, 'rgb'.index('g'), slicey, slicex] = self.GatR
+        weight[r, 'rgb'.index('b'), slicey, slicex] = self.BatR
+
+        # Finishing off blue bayer locations
+        b = bayer_pattern == 'rgb'.index('b')
+        slicey, slicex = kslice[rows[b]], kslice[cols[b]]
+        weight[b, 'rgb'.index('g'), slicey, slicex] = self.GatB
+        weight[b, 'rgb'.index('r'), slicey, slicex] = self.RatB
+
+        # greens get a bit more interesting because there are two
+        # types: one in red rows, and one in blue rows.
+        g, = where(bayer_pattern == 'rgb'.index('g'))
+        # read "gbr" as green pixel in blue row, red column
+        if any(b[:2]):  # if b is in the first row.  
+            gbr, grb = g
+        else:
+            grb, gbr = g
+        slicey, slicex = kslice[rows[grb]], kslice[cols[grb]]
+        weight[grb, 'rgb'.index('r'), slicey, slicex] = self.RatGRB
+        weight[grb, 'rgb'.index('b'), slicey, slicex] = self.BatGRB
+        slicey, slicex = kslice[rows[gbr]], kslice[cols[gbr]]
+        weight[gbr, 'rgb'.index('r'), slicey, slicex] = self.RatGBR
+        weight[gbr, 'rgb'.index('b'), slicey, slicex] = self.BatGBR
+
+        # apply YUV to RGB transform if necessary.  This is equivalent
+        # to scaling values AFTER applying filter.
+        for i, scale in enumerate(scales):
+            weight[:, i] *= float(scale)
+
+        # create the convulotion.
+        self.module = Conv2d(1, 12, (self.k+1, self.k+1), 2)
+        self.module.weight.data[:] = weight.reshape(12, 1, self.k+1, self.k+1)
+
+    def simple_init(self):
+        # generated by reading a 'demosaic.cpp' sent to me
+        self.GatR = tensor([[0, 1, 0],
+                            [1, 0, 1],
+                            [0, 1, 0]]
+                           ) / 4
+        # read "GRB" as green bayer location in red row, blue column.
+        self.RatGRB = tensor([[0, 0, 0],
+                              [1, 0, 1],
+                              [0, 0, 0]]
+                             ) / 2
+        self.RatB = tensor([[1, 0, 1],
+                            [0, 0, 0],
+                            [1, 0, 1]],
+                           ) / 4
+        self.k = 3
+        self.basic_init()
+
+    def malvar_init(self):
+        # kernels taken from https://www.microsoft.com/en-us/research/wp-content/uploads/2016/02/Demosaicing_ICASSP04.pdf
+        self.GatR = tensor([[ 0  , 0  ,-1  , 0  , 0  ],
+                            [ 0  , 0  , 2  , 0  , 0  ],
+                            [-1  , 2  , 4  , 2  ,-1  ],
+                            [ 0  , 0  , 2  , 0  , 0  ],
+                            [ 0  , 0  ,-1  , 0  , 0  ]],
+                           dtype=float32) / 8
+        # read "GRB" as green bayer location in red row, blue column.
+        self.RatGRB = tensor([[ 0  , 0  , 0.5, 0  , 0  ],
+                              [ 0  ,-1  , 0  ,-1  , 0  ],
+                              [-1  , 4  , 5  , 4  ,-1  ],
+                              [ 0  ,-1  , 0  ,-1  , 0  ],
+                              [ 0  , 0  , 0.5, 0  , 0  ]],
+                             dtype=float32) / 8
+        self.RatB = tensor([[ 0  , 0  ,-1.5, 0  , 0  ],
+                            [ 0  , 2  , 0  , 2  , 0  ],
+                            [-1.5, 0  , 6  , 0  ,-1.5],
+                            [ 0  , 2  , 0  , 2  , 0  ],
+                            [ 0  , 0  ,-1.5, 0  , 0  ]],
+                           dtype=float32) / 8
+        self.k = 5
+        self.basic_init()
+
+    def basic_init(self):
+        self.GatB = self.GatR
+        # read "GRB" as green bayer location in red row, blue column.
+        self.BatGBR = self.RatGRB
+        self.RatGBR = self.RatGRB.T
+        self.BatGRB = self.RatGBR
+        self.BatR = self.RatB
+
+
+def reshape_conv(old_conv):
+    assert (all(k in (3, 4) for k in old_conv.kernel_size)
+            and old_conv.stride == 2
+            and old_conv.in_channels == 3
+            and not old_conv.use_bias)
+    # first 2x2 is demosaic output, second 2x2 is this new conv's
+    # input.
+    weight = zeros(old_conv.out_channels, 2, 2, 3, 2, 2)
+    old_weight = old_conv.weight.data
+    # This is the best image I can use to try and describe (for 3x3
+    # starting kernel):
+    #
+    #     0   1   2
+    #   l       l       l
+    #   +---+---+---+ - + -
+    # 0 |rgb rgblrgb|rgb
+    #   +   +   +   +   +  0
+    # 1 |rgb rgblrgb|rgb
+    #   + - + - + - + - + -
+    # 2 |rgb rgblrgb|rgb
+    #   +---+---+---+   +  1
+    #   lrgb rgblrgb rgb
+    #   + - + - + - + - + -
+    #   l       l       l
+    #       0       1
+    #
+    # The left/top coordinates and ('|', '---') are in terms of the
+    # original kernel, and the right/bottom coordinates and ('l', ' -
+    # ') are in terms of the new input coordinates.  I use the
+    # coordinates above in the later comments.  The 3x3 box is the
+    # orignal conv kernel.  Each of the 4 2x2 blocks above have been
+    # transformed into one pixel of the demosaic output.
+
+    # (0, 0) in right/bottom coordinates
+    weight[:, 0, 0, :, 0, 0] = old_weight[:, :, 0, 0]
+    weight[:, 0, 1, :, 0, 0] = old_weight[:, :, 0, 1]
+    weight[:, 1, 0, :, 0, 0] = old_weight[:, :, 1, 0]
+    weight[:, 1, 1, :, 0, 0] = old_weight[:, :, 1, 1]
+    # (0, 1) in right/bottom coordinates
+    weight[:, 0, 0, :, 0, 1] = old_weight[:, :, 0, 2]
+    weight[:, 1, 0, :, 0, 1] = old_weight[:, :, 1, 2]
+    if old_conv.kernel_size[1] == 4:
+        weight[:, 0, 1, :, 0, 1] = old_weight[:, :, 0, 3]
+        weight[:, 1, 1, :, 0, 1] = old_weight[:, :, 1, 3]
+    # (1, 0) in right/bottom coordinates
+    weight[:, 0, 0, :, 1, 0] = old_weight[:, :, 2, 0]
+    weight[:, 0, 1, :, 1, 0] = old_weight[:, :, 2, 1]
+    if old_conv.kernel_size[0] == 4:
+        weight[:, 1, 0, :, 1, 0] = old_weight[:, :, 3, 0]
+        weight[:, 1, 1, :, 1, 0] = old_weight[:, :, 3, 1]
+    # (1, 1) in right/bottom coordinates
+    weight[:, 0, 0, :, 1, 1] = old_weight[:, :, 2, 2]
+    if old_conv.kernel_size[1] == 4:
+        weight[:, 0, 1, :, 1, 1] = old_weight[:, :, 2, 3]
+    if old_conv.kernel_size[0] == 4:
+        weight[:, 1, 0, :, 1, 1] = old_weight[:, :, 3, 2]
+    if all(k == 4 for k in old_conv.kernel_size):
+        weight[:, 1, 1, :, 1, 1] = old_weight[:, :, 3, 3]
+
+    conv = Conv2d(12, old_conv.out_channels, 2, 1,
+                  bias=old_conv.use_bias,
+                  padding=old_conv.padding == "same",
+                  groups=old_conv.groups)
+    conv.weight.data[:] = weight.reshape(old_conv.out_channels,
+                                         12, 2, 2)
+    return conv
+
+
+class UnfoldedDemosaic(Demosaic):
+    def forward(self, x):
+        x = self.module(x)
+        if askeras.use_keras:
+            return self.as_keras(x)
+        *B, C, H, W = x.shape
+        assert C == 12
+        permute = 2, 3, 0, 4, 1
+        permute = tuple(range(len(B))) + tuple(v + len(B) for v in permute)
+        return x.reshape(*B, 2, 2, 3, H, W
+                         ).permute(permute
+                                   ).reshape(*B, 3, 2 * H, 2 * W)
+
+    def clf(self, x):
+        *B, H, W, C = x.shape
+        permute = 2, 0, 1
+        permute = tuple(range(len(B))) + tuple(v + len(B) for v in permute)
+        x = self.module(tensor(x, dtype=float32).permute(permute))
+        permute = 3, 0, 4, 1, 2
+        permute = tuple(range(len(B))) + tuple(v + len(B) for v in permute)
+        return x.reshape(*B, 2, 2, 3, H // 2, W // 2
+                         ).permute(permute
+                                   ).reshape(*B, H, W, 3).detach().numpy()
+
+    def as_keras(self, x):
+        from keras.layers import Reshape, Permute
+        *_, H, W, _ = x.shape
+        return Reshape((H * 2, W * 2, 3))(
+            Permute((1, 3, 2, 4, 5))(
+                Reshape((H, W, 2, 2, 3))(x)
+            )
+        )

synet/katana.py

@@ -0,0 +1,4 @@
+"""katana.py includes imports which are compatible with Katana, and layer definitions that are only compatible with Katana.  However, Katana's capabilities are currently a subset of all other chip's capabilities, so it includes only imports for now."""
+
+from .layers import (Conv2dInvertedResidual, Head, SWSBiRNN, SRNN)
+from .base import askeras, Conv2d, Cat, ReLU, BatchNorm, Grayscale

synet/layers.py

@@ -0,0 +1,439 @@
+"""layers.py is the high level model building layer of synet.  It
+defines useful composite layers which are compatible with multiple
+chips.  Because it is built with layers from base.py, exports come
+"free".  As a rule of thumb to differentiate between base.py,
+layers.py:
+
+- base.py should only import from torch, keras, and tensorflow.
+- layers.py should only import from base.py.
+
+If you sublcass from something in base.py OTHER than Module, you
+should add a test case for it in tests/test_keras.py.
+
+"""
+from typing import Union, Tuple, Optional
+
+from .base import (ReLU, BatchNorm, Conv2d, Module, Cat, Sequential,
+                   RNN, GRU, LSTM, Transpose, Reshape, Flip, Add,
+                   Shape, ModuleList, ChannelSlice, DepthwiseConv2d)
+
+
+# because this module only reinterprets Conv2d parameters, the test
+# case is omitted.
+class DepthwiseConv2d(DepthwiseConv2d):
+    def __init__(self,
+                 channels: int,
+                 kernel_size: Union[int, Tuple[int, int]],
+                 stride: int = 1,
+                 bias: bool = False,
+                 padding: Optional[bool] = True):
+        super().__init__(channels, channels, kernel_size, stride,
+                         bias, padding, groups=channels)
+
+
+class InvertedResidual(Module):
+    """
+    Block of conv2D -> activation -> linear pointwise with residual concat.
+    Inspired by Inverted Residual blocks which are the main building block
+    of MobileNet.  It is stable and gives low peek memory before and after.
+    Additionally, the computations are extremely efficient on our chips
+    """
+
+    def __init__(self, in_channels, expansion_factor,
+                 out_channels=None, stride=1, kernel_size=3,
+                 skip=True):
+        """This inverted residual takes in_channels to
+        in_channels*expansion_factor with a 3x3 convolution.  Then
+        after a batchnorm and ReLU, the activations are taken back
+        down to in_channels (or out_channels, if specified).  If
+        out_channels is not specified (or equals in_channels), and the
+        stride is 1, then the input will be added to the output before
+        returning."""
+        super().__init__()
+        if out_channels is None:
+            out_channels = in_channels
+        hidden = int(in_channels * expansion_factor)
+        self.layers = Sequential(Conv2d(in_channels,
+                                        out_channels=hidden,
+                                        kernel_size=kernel_size,
+                                        stride=stride),
+                                 BatchNorm(hidden),
+                                 ReLU(6),
+                                 Conv2d(in_channels=hidden,
+                                        out_channels=out_channels,
+                                        kernel_size=1),
+                                 BatchNorm(out_channels))
+        self.stride = stride
+        self.cheq = in_channels == out_channels and skip
+
+    def forward(self, x):
+        y = self.layers(x)
+        if self.stride == 1 and self.cheq:
+            return x + y
+        return y
+
+
+# for backwards compatibility
+Conv2dInvertedResidual = InvertedResidual
+
+
+class Head(Module):
+    def __init__(self, in_channels, out_channels, num=4):
+        """Creates a sequence of convolutions with ReLU(6)'s.
+in_channels features are converted to out_channels in the first
+convolution.  All other convolutions have out_channels going in and
+out of that layer.  num (default 4) convolutions are used in total.
+
+        """
+        super().__init__()
+        self.relu = ReLU(6)
+        out_channels = [in_channels] * (num - 1) + [out_channels]
+        self.model = Sequential(*(Sequential(Conv2d(in_channels,
+                                                    out_channels,
+                                                    3, bias=True),
+                                             self.relu)
+                                  for out_channels in out_channels))
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class CoBNRLU(Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3,
+                 stride=1, bias=False, padding=True, groups=1,
+                 max_val=6, name=None):
+        super().__init__()
+        self.module = Sequential(Conv2d(in_channels, out_channels,
+                                        kernel_size, stride, bias, padding,
+                                        groups),
+                                 BatchNorm(out_channels),
+                                 ReLU(max_val, name=name))
+
+    def forward(self, x):
+        return self.module(x)
+
+
+class GenericSRNN(Module):
+    """
+    Implements GenericSRNN (Generic Separable RNN), which processes an input tensor sequentially
+    along its X-axis and Y-axis using two RNNs.
+    This approach first applies an RNN along the X-axis of the input, then feeds the resulting tensor
+    into another RNN along the Y-axis.
+
+    Args:
+        - hidden_size_x (int): The number of features in the hidden state of the X-axis RNN.
+        - hidden_size_y (int): The number of features in the hidden state of the Y-axis RNN,
+          which also determines the output size.
+        - num_layers (int, optional): Number of recurrent layers for each RNN, defaulting to 1.
+
+    Returns:
+        - output (tensor): The output tensor from the Y-axis RNN.
+        - hn_x (tensor): The final hidden state from the X-axis RNN.
+        - hn_y (tensor): The final hidden state from the Y-axis RNN.
+    """
+
+    def __init__(self, hidden_size_x: int, hidden_size_y: int) -> None:
+        super(GenericSRNN, self).__init__()
+        self.output_size_x = hidden_size_x
+        self.output_size_y = hidden_size_y
+
+        self.transpose = Transpose()
+        self.reshape = Reshape()
+        self.get_shape = Shape()
+
+    def forward(self, x, rnn_x: Module, rnn_y: Module):
+        """
+        Performs the forward pass for the HierarchicalRNN module.
+
+        Args:
+            - x (tensor): Input tensor with shape [batch_size, channels, height, width]
+
+        Returns:
+            - output (tensor): The final output tensor from the Y-axis RNN.
+        """
+        batch_size, channels, H, W = self.get_shape(x)
+
+        # Rearrange the tensor to (batch_size*H, W, channels) for
+        # RNN processing over height This step prepares the data by aligning
+        # it along the width, treating each row separately.
+        # the keep_channel_last is true in case using channel last, but the
+        # assumption of the transpose dims is that the input is channels first,
+        # so is sign to keep the channels in the last dim.
+        x_w = self.transpose(x, (0, 2, 3, 1),
+                             keep_channel_last=True)  # Rearranges to (batch_size, H, W, channels)
+        x_w = self.reshape(x_w, (batch_size * H, W, channels))
+
+        output_x, _ = rnn_x(x_w)
+
+        # Prepare the output from the X-axis RNN for Y-axis processing by
+        # rearranging it to (batch_size*W, H, output_size_x),
+        # enabling RNN application over width.
+        output_x_reshape = self.reshape(output_x,
+                                        (batch_size, H, W, self.output_size_x))
+        output_x_permute = self.transpose(output_x_reshape, (0, 2, 1, 3))
+        output_x_permute = self.reshape(output_x_permute,
+                                        (batch_size * W, H, self.output_size_x))
+
+        output, _ = rnn_y(output_x_permute)
+
+        # Reshape and rearrange the final output to
+        # (batch_size, channels, height, width),
+        # restoring the original input dimensions with the transformed data.
+        output = self.reshape(output, (batch_size, W, H, self.output_size_y))
+        output = self.transpose(output, (0, 3, 2, 1), keep_channel_last=True)
+
+        return output
+
+
+class SRNN(GenericSRNN):
+    """
+    Implements the Separable Recurrent Neural Network (SRNN).
+    This model extends a standard RNN by introducing separability in processing.
+    """
+
+    def __init__(self, input_size: int, hidden_size_x: int, hidden_size_y: int,
+                 base: str = 'RNN', num_layers: int = 1, bias: bool = True,
+                 batch_first: bool = True, dropout: float = 0.0,
+                 bidirectional: bool = False) -> None:
+        """
+        Initializes the SRNN model with the given parameters.
+
+        Parameters:
+        - input_size: The number of expected features in the input `x`
+        - hidden_size_x: The number of features in the hidden state `x`
+        - hidden_size_y: The number of features in the hidden state `y`
+        - base: The type of RNN to use (e.g., 'RNN', 'LSTM', 'GRU')
+        - num_layers: Number of recurrent layers. E.g., setting `num_layers=2`
+        would mean stacking two RNNs together
+        - bias: If `False`, then the layer does not use bias weights `b_ih` and
+        `b_hh`. Default: `True`
+        - batch_first: If `True`, then the input and output tensors are provided
+        as (batch, seq, feature). Default: `True`
+        - dropout: If non-zero, introduces a `Dropout` layer on the outputs of
+        each RNN layer except the last layer,
+                   with dropout probability equal to `dropout`. Default: 0
+       - bidirectional: If `True`, becomes a torch implementation of
+       bidirectional RNN (two RNN blocks, one for the forward pass and one
+       for the backward). Default: `False`
+
+        Creates two `RNN` instances for processing in `x` and `y`
+        dimensions, respectively.
+
+        From our experiments, we found that the best results were
+        obtained with the following parameters:
+        base='RNN', num_layers=1, bias=True, batch_first=True, dropout=0
+        """
+        super(SRNN, self).__init__(hidden_size_x, hidden_size_y)
+
+        # Dictionary mapping base types to their respective PyTorch class
+        RNN_bases = {'RNN': RNN,
+                     'GRU': GRU,
+                     'LSTM': LSTM}
+
+        self.rnn_x = RNN_bases[base](input_size=input_size,
+                                     hidden_size=hidden_size_x,
+                                     num_layers=num_layers,
+                                     bias=bias,
+                                     batch_first=batch_first,
+                                     dropout=dropout,
+                                     bidirectional=bidirectional)
+
+        self.rnn_y = RNN_bases[base](input_size=hidden_size_x,
+                                     hidden_size=hidden_size_y,
+                                     num_layers=num_layers,
+                                     bias=bias,
+                                     batch_first=batch_first,
+                                     dropout=dropout,
+                                     bidirectional=bidirectional)
+
+        # Output sizes of the model in the `x` and `y` dimensions.
+        self.output_size_x = hidden_size_x
+        self.output_size_y = hidden_size_y
+
+    def forward(self, x):
+        """
+        Defines the forward pass of the SRNN.
+
+        Parameters:
+        - x: The input tensor to the RNN
+
+        Returns:
+        - The output of the SRNN after processing the input tensor
+        `x` through both the `x` and `y` RNNs.
+        """
+        output = super().forward(x, self.rnn_x, self.rnn_y)
+        return output
+
+
+class SWSBiRNN(GenericSRNN):
+    """
+    Implements the Weights Shared Bi-directional Separable Recurrent Neural Network (WSBiSRNN).
+    This model extends a standard RNN by introducing bi-directionality and separability in processing,
+    with weight sharing to reduce the number of parameters.
+    """
+
+    def __init__(self, input_size: int, hidden_size_x: int, hidden_size_y: int,
+                 base: str = 'RNN', num_layers: int = 1, bias: bool = True,
+                 batch_first: bool = True, dropout: float = 0.0) -> None:
+        """
+        Initializes the WSBiSRNN model with the given parameters.
+
+        Parameters:
+        - input_size: The number of expected features in the input `x`
+        - hidden_size_x: The number of features in the hidden state `x`
+        - hidden_size_y: The number of features in the hidden state `y`
+        - base: The type of RNN to use (e.g., 'RNN', 'LSTM', 'GRU')
+        - num_layers: Number of recurrent layers. E.g., setting `num_layers=2`
+        would mean stacking two RNNs together
+        - bias: If `False`, then the layer does not use bias weights `b_ih` and
+        `b_hh`. Default: `True`
+        - batch_first: If `True`, then the input and output tensors are provided
+        as (batch, seq, feature). Default: `True`
+        - dropout: If non-zero, introduces a `Dropout` layer on the outputs of
+        each RNN layer except the last layer,
+                   with dropout probability equal to `dropout`. Default: 0
+
+        Creates two `WSBiRNN` instances for processing in `x` and `y`
+        dimensions, respectively.
+
+        From our experiments, we found that the best results were
+        obtained with the following parameters:
+        base='RNN', num_layers=1, bias=True, batch_first=True, dropout=0
+        """
+        super(SWSBiRNN, self).__init__(hidden_size_x, hidden_size_y)
+
+        self.rnn_x = WSBiRNN(input_size=input_size,
+                             hidden_size=hidden_size_x,
+                             num_layers=num_layers,
+                             base=base,
+                             bias=bias,
+                             batch_first=batch_first,
+                             dropout=dropout)
+
+        self.rnn_y = WSBiRNN(input_size=hidden_size_x,
+                             hidden_size=hidden_size_y,
+                             num_layers=num_layers,
+                             base=base,
+                             bias=bias,
+                             batch_first=batch_first,
+                             dropout=dropout)
+
+        # Output sizes of the model in the `x` and `y` dimensions.
+        self.output_size_x = hidden_size_x
+        self.output_size_y = hidden_size_y
+
+    def forward(self, x):
+        """
+        Defines the forward pass of the WSBiSRNN.
+
+        Parameters:
+        - x: The input tensor to the RNN
+
+        Returns:
+        - The output of the WSBiSRNN after processing the input tensor `x` through both the `x` and `y` RNNs.
+        """
+        output = super().forward(x, self.rnn_x, self.rnn_y)
+        return output
+
+
+class WSBiRNN(Module):
+    """
+    WSBiRNN (Weight-Shared Bidirectional) RNN is a custom implementation of a bidirectional
+    RNN that processes input sequences in both forward and reverse directions
+    and combines the outputs. This class manually implements
+    bidirectional functionality using a specified base RNN (e.g., vanilla RNN, GRU, LSTM)
+    and combines the forward and reverse outputs.
+
+    Attributes:
+        rnn (Module): The RNN module used for processing sequences in the forward direction.
+        hidden_size (int): The size of the hidden layer in the RNN.
+        flip (Flip): An instance of the Flip class for reversing the sequence order.
+        add (Add): An instance of the Add class for combining forward and reverse outputs.
+    """
+
+    def __init__(self, input_size: int, hidden_size: int, num_layers: int = 1,
+                 base: str = 'RNN', bias: bool = True, batch_first: bool = True,
+                 dropout: float = 0.0) -> None:
+        """
+        Initializes the BiDirectionalRNN module with the specified parameters.
+
+        Parameters:
+            input_size (int): The number of expected features in the input `x`.
+            hidden_size (int): The number of features in the hidden state `h`.
+            num_layers (int, optional): Number of recurrent layers. Default: 1.
+            base (str, optional): Type of RNN ('RNN', 'GRU', 'LSTM'). Default: 'RNN'.
+            bias (bool, optional): If False, then the layer does not use bias weights. Default: True.
+            batch_first (bool, optional): If True, then the input and output tensors are provided
+                                          as (batch, seq, feature). Default: True.
+            dropout (float, optional): If non-zero, introduces a Dropout layer on the outputs of
+                                       each RNN layer except the last layer. Default: 0.
+        """
+        super(WSBiRNN, self).__init__()
+
+        # Dictionary mapping base types to their respective PyTorch class
+        RNN_bases = {'RNN': RNN,
+                     'GRU': GRU,
+                     'LSTM': LSTM}
+
+        # Initialize the forward RNN module
+        self.rnn = RNN_bases[base](input_size=input_size,
+                                   hidden_size=hidden_size,
+                                   num_layers=num_layers,
+                                   bias=bias,
+                                   batch_first=batch_first,
+                                   dropout=dropout,
+                                   bidirectional=False)
+
+        # Initialize utilities for flipping sequences and combining outputs
+        self.flip = Flip()
+        self.add = Add()
+
+    def forward(self, x):
+        """
+        Defines the forward pass for the bidirectional RNN.
+
+        Parameters:
+            x (Tensor): The input sequence tensor.
+
+        Returns:
+            Tensor: The combined output of the forward and reverse processed sequences.
+            _: Placeholder for compatibility with the expected RNN output format.
+        """
+        # Reverse the sequence for processing in the reverse direction
+        x_reverse = self.flip(x, [1])
+
+        # Process sequences in forward and reverse directions
+        out_forward, _ = self.rnn(x)
+        out_reverse, _ = self.rnn(x_reverse)
+
+        # Flip the output from the reverse direction to align with forward
+        # direction
+        out_reverse_flip = self.flip(out_reverse, [1])
+
+        # Combine the outputs from the forward and reverse directions
+        return self.add(out_reverse_flip, out_forward), _
+
+
+class S2f(Module):
+    # Synaptics C2f.  Constructed from tflite inspection
+    def __init__(self, in_channels, n, out_channels=None, nslice=None,
+                 skip=True):
+        super().__init__()
+        if out_channels is None:
+            out_channels = in_channels
+        if nslice is not None:
+            c = nslice
+        else:
+            c = in_channels // 2
+        self.slice = ChannelSlice(slice(c))
+        self.ir = ModuleList([InvertedResidual(c, 1, skip=skip is True)
+                              for _ in range(n)])
+        self.cat = Cat()
+        self.decode = CoBNRLU(in_channels + n*c, out_channels, bias=True)
+
+    def forward(self, x):
+        out = [x]
+        y = self.slice(x)
+        for ir in self.ir:
+            out.append(y := ir(y))
+        return self.decode(self.cat(out))

synet/legacy.py

@@ -0,0 +1,151 @@
+from numpy import array
+
+from os.path import join, dirname
+from json import load
+from tensorflow.keras.models import load_model
+def get_katananet_model(model_path, input_shape, low_thld, **kwds):
+    """Load katananet model.
+
+model_dir: str
+    path to directory with model.h5.
+input_shape: iterable of ints
+    shape of the cell run.
+
+    """
+    raw_model = load_model(model_path, compile=False)
+    anchor_params = load(open(join(dirname(model_path), "anchors.json")))
+    anchors = gen_anchors(input_shape, **anchor_params)
+    def model(image):
+        (deltas,), (scores,) = raw_model.predict_on_batch(preproc(image))
+        # low thld
+        keep = scores.max(1) > low_thld
+        deltas, anchor_keep, scores = deltas[keep], anchors[keep], scores[keep]
+        # get_abs_coords only get coordinates relative to cell
+        boxes = get_abs_coords(deltas, anchor_keep,
+                               training_scale=.2, training_shift=0,
+                               maxx=image.shape[-1], maxy=image.shape[-2])
+        # apply nms
+        boxes, scores = nms(boxes, scores, threshold=.3)
+        return boxes, scores.squeeze(-1)
+
+    return model
+
+from numpy import float32, expand_dims
+def preproc(image):
+    """Convert image values from integer range [0,255] to float32
+    range [-1,1)."""
+    if len(image.shape) < 3:
+        image = expand_dims(image, 0)
+    return image.astype(float32) / 128 - 1
+
+from numpy import zeros, arange, concatenate
+from math import ceil
+def gen_anchors(image_shape, strides, sizes, ratios, scales):
+    imy, imx = image_shape
+    all_anchors = []
+    scales = array(scales).reshape(-1, 1)
+    ratios = array(ratios).reshape(-1, 1, 1)**.5
+    for stride, size in zip(strides, sizes):
+        py, px = ceil(imy/stride), ceil(imx/stride)
+        anchors = zeros((py, px, len(ratios), len(scales), 4))
+        # anchors as (xc, yc, w, h)
+        anchors[...,2:] = size * scales
+        # apply ratios
+        anchors[...,2] /= ratios[...,0]
+        anchors[...,3] *= ratios[...,0]
+        # convert to xyxy
+        anchors[...,:2] -= anchors[...,2:]/2
+        anchors[...,2:] /= 2
+        # add offsets for xy position
+        anchors[...,0::2] += ((arange(px) + 0.5) * stride).reshape(-1,1,1,1)
+        anchors[...,1::2] += ((arange(py) + 0.5) * stride).reshape(-1,1,1,1,1)
+        all_anchors.append(anchors.reshape(-1, 4))
+    return concatenate(all_anchors)
+
+from numpy import clip, newaxis
+def get_abs_coords(deltas, anchors, training_scale, training_shift,
+                   maxx, maxy):
+    """Convert model output (deltas) into "absolute" coordinates.
+    Note: absolute coordinates here are still relative to the grid
+    cell being run.
+
+deltas: ndarray
+    nx4 array of xyxy values.
+anchors: ndarray
+    nx4 array of ofsets.
+training_scale: float
+    scale specific to our training code.  For us always set to .2.
+training_shift: float
+    shift specific to our training code.  For us is always 0.
+maxx: float
+    Max x value. Used to clip final results to fit in cell.
+maxy: float
+    Max y value. Used to clip final results to fit in cell.
+
+    """
+    width, height = (anchors[:, 2:4] - anchors[:, 0:2]).T
+    deltas = deltas * training_scale + training_shift
+    deltas[:,0::2] *= width [...,newaxis]
+    deltas[:,1::2] *= height[...,newaxis]
+    boxes = deltas + anchors
+    boxes[:, 0::2] = clip(boxes[:, 0::2], 0, maxx)
+    boxes[:, 1::2] = clip(boxes[:, 1::2], 0, maxy)
+    return boxes
+
+from numpy import argsort, maximum, minimum
+def nms(boxes, score, threshold):
+    """
+    Non-maxima supression to remove redundant boxes
+    :param bounding_boxes: Input box coordinates
+    :param confidence_score: Confidence scores for each box
+    :param labels: Class label for each box
+    :param threshold: Only boxes above this threshold are selected
+    :return:
+    Final detected boxes
+    """
+    if not len(boxes):
+        return boxes, score
+
+    # coordinates of bounding boxes
+    all_x1 = boxes[:, 0]
+    all_y1 = boxes[:, 1]
+    all_x2 = boxes[:, 2]
+    all_y2 = boxes[:, 3]
+
+    # Picked bounding boxes
+    picked_boxes = []
+    picked_score = []
+
+    # Compute areas of bounding boxes
+    areas = (all_y2 - all_y1 + 1) * (all_x2 - all_x1 + 1)
+
+    # Sort by confidence score of bounding boxes
+    order = argsort(-score.max(-1))
+
+    # Iterate bounding boxes
+    while order.size > 0:
+        # The index of largest confidence score
+        index = order[0]
+        order = order[1:]
+
+        # Pick the bounding box with largest confidence score
+        picked_boxes.append(boxes[index])
+        picked_score.append(score[index])
+
+        # Compute ordinates of intersection-over-union(IOU)
+        y1 = maximum(all_y1[index], all_y1[order])
+        x1 = maximum(all_x1[index], all_x1[order])
+        y2 = minimum(all_y2[index], all_y2[order])
+        x2 = minimum(all_x2[index], all_x2[order])
+
+        # Compute areas of intersection-over-union
+        w = maximum(0.0, x2 - x1 + 1)
+        h = maximum(0.0, y2 - y1 + 1)
+        intersection = w * h
+
+        # Compute the ratio between intersection and union
+        ratio = intersection / (areas[index] + areas[order] - intersection)
+
+        order = order[ratio < threshold]
+
+    return array(picked_boxes), array(picked_score)

synet/metrics.py

@@ -0,0 +1,328 @@
+from argparse import ArgumentParser, Namespace
+from glob import glob
+from os import listdir, makedirs
+from os.path import join, basename, isfile, isdir, isabs
+
+from numpy import genfromtxt
+from torch import tensor, stack, cat, empty
+from ultralytics.utils.ops import xywh2xyxy
+from ultralytics.data.utils import check_det_dataset, img2label_paths
+
+
+aP_curve_points = 10000
+
+
+def parse_opt() -> Namespace:
+    parser = ArgumentParser()
+    parser.add_argument("data_yamls", nargs="+")
+    parser.add_argument("--out-dirs", nargs="+")
+    parser.add_argument("--project")
+    parser.add_argument("--name")
+    parser.add_argument("--print-jobs", action="store_true")
+    parser.add_argument("--precisions", nargs="+", type=float, required=True,
+                        help="CANNOT BE SPECIFIED WITH --precisions=...' "
+                        "SYNTAX: MUST BE '--precisions PREC1 PREC2 ...'")
+    return parser.parse_known_args()[0]
+
+
+def txt2xyxy(txt : str, conf=False) -> tensor:
+    """Convert txt path to array of (cls, x1, y1, x2, y2[, conf])"""
+    a = tensor(genfromtxt(txt, ndmin=2))
+    if not len(a):
+        return empty(0, 5+int(conf))
+    a[:, 1:5] = xywh2xyxy(a[:, 1:5])
+    return a
+
+
+def get_gt(data_yaml : str) -> dict:
+    """Obtain {"###.txt" : (Mx5 array)} dictionary mapping each data
+sample to an array of ground truths.  See get_pred().
+
+    """
+    cfg = check_det_dataset(data_yaml)
+    path = cfg.get('test', cfg['val'])
+    f = []
+    for p in path if isinstance(path, list) else [path]:
+        if isdir(p):
+            f += glob(join(p, "**", "*.*"), recursive=True)
+        else:
+            f += [t if isabs(t) else join(dirname(p), t)
+                  for t in open(p).read().splitlines()]
+    print('getting gt')
+    return {basename(l): txt2xyxy(l, conf=False)
+            for l in img2label_paths(f)}
+
+
+def get_pred(pred : str) -> dict:
+    """from model output dir from validation, pred, obtain {"###.txt"
+: (Mx6 array)} mapping each data sample to array of predictions (with
+confidence) on that sample.  See get_gt().
+
+    """
+    print('getting pred')
+    return {name : txt2xyxy(join(pred, name), conf=True)
+            for name in listdir(pred)}
+
+
+from yolov5.val import process_batch
+# from torchvision.ops import box_iou
+# def validate_preds(sample_pred, sample_gt):
+#     box_iou(sample_pred)
+#     correct = zeros(len(sample_pred)).astype(bool)
+
+def get_tp_ngt(gt : dict, pred : dict) -> tuple:
+    """From ground truth and prediction dictionaries (as given by
+get_gt() and get_pred() funcs resp.), generate a single Mx3 array, tp,
+for the entire dataset, as well as a dictionary, gt = { (class : int)
+: (count : int) }, giving the count of each ground truth.  The array
+is interpreted as meaning there are M predictions denoted by (conf,
+class, TP) giving the network predicted confidence, the network
+predicted class, and a flag TP which is 1 if the sample is considered
+a true positive.
+
+    """
+    # after this point, we don't care about which pred came from which
+    # data sample in this data split
+    tp = cat([stack((pred[fname][:, 5], # conf
+                     pred[fname][:, 0], # class
+                     process_batch(pred[fname][:,[1,2,3,4,5,0]],
+                                   gt[fname],
+                                   tensor([.5])
+                                   ).squeeze(1)), # TP
+                    -1)
+              for fname in pred])
+    l = cat([gt[fname][:, 0] for fname in pred])
+    ngt = {int(c.item()) : (l == c).sum() for c in l.unique()}
+    return tp, ngt
+
+from torch import cumsum, arange, linspace
+from numpy import interp
+from matplotlib.pyplot import (plot, legend, title, xlabel, ylabel,
+                               savefig, clf, scatter, grid, xlim, ylim)
+def get_aps(tp : tensor, ngt : dict, precisions : list, label : str,
+            project : str, glob_confs : [list, None] = None) -> list:
+    """This is the main metrics AND plotting function.  All other
+functions exist to "wrangle" the data into an optimal format for this
+function. From a 'tp' tensor and 'ngt' dict (see get_tp_ngt()),
+compute various metrics, including the operating point at
+'precisions'[c] for each class c.  Plots are labeled and nammed based
+on 'label', and placed in the output dir 'project'.  Additionally, if
+glob_confs is also given, plot the operating point at that confidence
+threshold.  Returns the confidence threshold corresponding to each
+precision threshold in 'precisions'.
+
+    """
+    # if there are fewer precision thresholds specified than classes
+    # present, and only one precision is specified, use that precision
+    # for all classes
+    if max(ngt) > len(precisions) - 1:
+        if len(precisions) == 1:
+            print("applying same precision to all classes")
+            precisions *= max(ngt)
+        else:
+            print("specified", len(precisions), "precisions, but have",
+                  max(ngt)+1, "classes")
+            exit()
+    # Main loop.  One for each class.  AP calculated at the end
+    AP, confs, op_P, op_R, half_P, half_R = [], [], [], [], [], []
+    if glob_confs is not None: glob_P, glob_R = [], []
+    for cls, prec in enumerate(precisions):
+        print("For class:", cls)
+
+        # choose class and omit class field
+        selected = tp[tp[:,1] == cls][:,::2]
+
+        # sort descending
+        selected = selected[selected[:, 0].argsort(descending=True)]
+
+        # calculate PR values
+        assert len(selected.shape) == 2
+        tpcount = cumsum(selected[:,1], 0).numpy()
+        P = tpcount / arange(1, len(tpcount) + 1)
+        R = tpcount / ngt.get(cls, 0)
+        # enforce that P should be monotone
+        P = P.flip(0).cummax(0)[0].flip(0)
+
+        # calculate operating point from precision.
+        # operating index is where the precision last surpasses precision thld
+        # argmax on bool array returns first time condition is met.
+        # Precision is not monotone, so need to reverse, argmax, then find ind
+        assert len(P.shape) == 1
+        confs.append(selected[(P < prec).byte().argmax() -1, 0])
+        op_ind = (selected[:,0] <= confs[-1]).byte().argmax() - 1
+        op_P.append(P[op_ind])
+        op_R.append(R[op_ind])
+        print(f"Conf, Precision, Recall at operating point precision={prec}")
+        print(f"{confs[-1]:.6f}, {op_P[-1]:.6f}, {op_R[-1]:.6f}")
+
+        if glob_confs is not None:
+            # if glob threshold is passed, also find that PR point
+            glob_ind = (selected[:,0] <= glob_confs[cls]).byte().argmax() - 1
+            glob_P.append(P[glob_ind])
+            glob_R.append(R[glob_ind])
+            print("Conf, Precision, Recall at global operating point:")
+            print(f"""{glob_confs[cls]:.6f}, {glob_P[-1]
+                      :.6f}, {glob_R[-1]:.6f}""")
+
+        # show .5 conf operating point
+        half_ind = (selected[:,0] <= .5).byte().argmax() - 1
+        half_P.append(P[half_ind])
+        half_R.append(R[half_ind])
+        print(f"Conf, Precision, Recall at C=.5 point")
+        print(f"{.5:.6f}, {half_P[-1]:.6f}, {half_R[-1]:.6f}")
+
+        # generate plotting points/AP calc points
+        Ri = linspace(0, 1, aP_curve_points)
+        Pi = interp(Ri, R, P)
+        # use these values for AP calc over raw to avoid machine error
+        AP.append(Pi.sum() / aP_curve_points)
+        print("class AP:", AP[-1].item(), end="\n\n")
+        plot(Ri, Pi, label=f"{cls}: AP={AP[-1]:.6f}")
+
+    # calculate mAP
+    mAP = sum(AP)/len(AP)
+    print("mAP:", mAP, end="\n\n\n")
+    title(f"{basename(label)} mAP={mAP:.6f}")
+
+    # plot other points
+    scatter(op_R, op_P, label="precision operating point")
+    scatter(half_R, half_P, label=".5 conf")
+    if glob_confs is not None:
+        scatter(glob_R, glob_P, label="global operating point")
+
+    # save plot
+    legend()
+    xlabel("Recall")
+    ylabel("Precision")
+    grid()
+    xlim(0, 1)
+    ylim(0, 1)
+    savefig(join(project, f"{basename(label)}.png"))
+    clf()
+
+    return confs
+
+def metrics(data_yamls : list, out_dirs : list, precisions : list,
+            project : str) -> None:
+    """High level function for computing metrics and generating plots
+for the combined data plus each data split.  Requires list of data
+yamls, data_yamls, model output dirs, out_dirs, classwise precision
+thresholds, precisions, and output dir, project.
+
+    """
+    tp_ngt = {}
+    for data_yaml, out_dir in zip(data_yamls, out_dirs):
+        tp_ngt[data_yaml] = get_tp_ngt(get_gt(data_yaml),
+                                       get_pred(join(out_dir, 'labels')))
+    print("Done reading results.  Results across all data yamls:", end="\n\n")
+    confs = get_aps(cat([tp for tp, _ in tp_ngt.values()]),
+                    {c : sum(ngt.get(c, 0) for _, ngt in tp_ngt.values())
+                     for c in set.union(*(set(ngt.keys())
+                                          for _, ngt in tp_ngt.values()))},
+                    precisions,
+                    "all",
+                    project)
+    if len(tp_ngt) == 1:
+        return
+    for data_yaml, (tp, ngt) in tp_ngt.items():
+        print("Results for", data_yaml, end="\n\n")
+        get_aps(tp, ngt, precisions, data_yaml, project, confs)
+
+from sys import argv
+from synet.__main__ import main as synet_main
+def run(data_yamls : list, out_dirs : [list, None], print_jobs : bool,
+        precisions : list, project : [str, None], name : None):
+    """Entrypoint function.  Compute metrics of model on data_yamls.
+
+If out_dirs is specified, it should be a list of output directories
+used for validation runs on the datasets specified by data_yamls (in
+the same order).
+
+If out_dirs is not specified, then all necessary validation args
+should be specified in command-line args (sys.argv).  In this case, it
+will run validation of your model on each specified data yaml before
+attempting to compute metrics.
+
+If print_jobs is specified, then the commands to run the various
+validation jobs are printed instead.  This is useful if you would like
+to run the validation jobs in parallel.
+
+If project is specified, this will be used as the base output
+directory for plots and generated validation jobs.
+
+name should never be specified.  validation job names are generated by
+this function, so you must not try to specify your own.
+
+precisions is a list of precision thresholds.  This is used as an
+operating point which is also reported by the metrics here.  It is
+either one value (used for all classes), or a list of values
+correspoinding to the labels in order.
+
+    """
+
+    # decide output dir
+    assert name is None, "--name specified by metrics.  Do not specify"
+    makedirs(project, exist_ok=True)
+    if project is None:
+        project = "metrics"
+        argv.append(f"--project={project}")
+
+    # if val was already run, just compute metrics
+    if out_dirs is not None:
+        assert len(out_dirs) == len(data_yamls), \
+            "Please specify one output for each data yaml"
+        print("Using prerun results to compute metrcis")
+        return metrics(data_yamls, out_dirs, precisions, project)
+
+    ## modify argv
+    # add necessary flags
+    for flag in "--save-conf", '--save-txt', '--exist-ok':
+        if flag not in argv:
+            argv.append(flag)
+    # remove precisions flag from args
+    argv.remove("--precisions")
+    if print_jobs:
+        argv.remove("--print-jobs")
+    rm = [arg for precison in precisions for arg in argv
+          if arg.isnumeric() and -.0001 <= float(arg) - precision <= .0001]
+    for r in rm:
+        if r in argv: argv.remove(r)
+    # remove data yamls from args
+    for data_yaml in data_yamls: argv.remove(data_yaml)
+    # run validation
+    argv.insert(1, "val")
+
+    ## generate val jobs
+    if print_jobs:
+        print("Submit the following jobs:")
+    out_dirs = []
+    for i, data_yaml in enumerate(data_yamls):
+        # specify data and out dir.
+        flags = f"--data={data_yaml}", f"--name=data-split{i}"
+        argv.extend(flags)
+        # run/print the job
+        print(" ".join(argv))
+        if not print_jobs:
+            print("starting job")
+            synet_main()
+            # main removes job type from argv, re-add it
+            argv.insert(1, "val")
+        # revert argv for next job
+        for flag in flags:
+            argv.remove(flag)
+        # keep track of output dirs in order
+        out_dirs.append(join(project, f"data-split{i}"))
+
+    ## calculate metrics
+    if print_jobs:
+        print("Once jobs finish, run:")
+
+    print(" ".join([argv[0], "metrics", *data_yamls, "--out-dirs",
+                    *out_dirs, "--precisions", *(str(prec)
+                                                 for prec in precisions)]))
+    if not print_jobs:
+        print("computing metrics")
+        return metrics(data_yamls, out_dirs, precisions, project)
+
+def main():
+    run(**vars(parse_opt()))

synet/quantize.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python
+from argparse import ArgumentParser
+from glob import glob
+from os.path import dirname, isabs, isdir, join, splitext
+from random import shuffle
+
+from cv2 import imread, resize
+from keras import Input, Model
+from numpy import float32
+from numpy.random import rand
+from tensorflow import int8, lite
+from torch import no_grad
+
+from .base import askeras
+from .backends import get_backend
+
+
+def parse_opt(args=None):
+    """parse_opt() is used to make it compatible with how yolov5
+obtains arguments.
+
+    """
+    parser = ArgumentParser()
+    parser.add_argument("--backend", type=get_backend,
+                        default=get_backend('ultralytics'))
+    parser.add_argument("--model", "--cfg", '--weights')
+    parser.add_argument("--image-shape", nargs=2, type=int)
+    parser.add_argument("--data")
+    parser.add_argument("--kwds", nargs="+", default=[])
+    parser.add_argument("--channels", "-c", default=3, type=int)
+    parser.add_argument("--number", "-n", default=500, type=int)
+    parser.add_argument("--val-post",
+                        help="path to sample image to validate on.")
+    parser.add_argument("--tflite",
+                        help="path to existing tflite (for validating).")
+    return parser.parse_args(args=args)
+
+
+def run(backend, image_shape, model, data, number, channels, kwds,
+        val_post, tflite):
+    """Entrypoint to quantize.py.  Quantize the model specified by
+weights (falling back to cfg), using samples from the data yaml with
+image shape image_shape, using only number samples.
+
+    """
+    backend.patch()
+    model = backend.maybe_grab_from_zoo(model)
+
+    if tflite is None:
+        tflite = get_tflite(backend, image_shape, model, data,
+                            number, channels, kwds)
+
+    if val_post:
+        backend.val_post(model, tflite, val_post, image_shape=image_shape)
+
+
+def get_tflite(backend, image_shape, model_path, data, number,
+               channels, kwds):
+
+    # maybe get image shape
+    if image_shape is None:
+        image_shape = backend.get_shape(model_path)
+
+    # generate keras model
+    ptmodel = backend.get_model(model_path)
+    inp = Input(image_shape+[channels], batch_size=1)
+    with askeras(imgsz=image_shape, quant_export=True,
+                 **dict(s.split("=") for s in kwds)), \
+         no_grad():
+        kmodel = Model(inp, ptmodel(inp))
+
+    print('model params:', kmodel.count_params())
+
+    # quantize the model
+    return quantize(kmodel, data, image_shape, number,
+                    splitext(model_path)[0]+".tflite",
+                    channels, backend=backend)
+
+
+def quantize(kmodel, data, image_shape, N=500, out_path=None, channels=1,
+             generator=None, backend=None):
+    """Given a keras model, kmodel, and data yaml at data, quantize
+using N samples reshaped to image_shape and place the output model at
+out_path.
+
+    """
+    # more or less boilerplate code
+    converter = lite.TFLiteConverter.from_keras_model(kmodel)
+    converter.optimizations = [lite.Optimize.DEFAULT]
+    converter.inference_input_type = int8
+    converter.inference_output_type = int8
+
+    if generator:
+        converter.representative_dataset = generator
+    elif data is None:
+        converter.representative_dataset = \
+            lambda: phony_data(image_shape, channels)
+    else:
+        converter.representative_dataset = \
+            lambda: representative_data(backend.get_data(data), image_shape, N, channels)
+
+    # quantize
+    tflite_quant_model = converter.convert()
+
+    # write out tflite
+    if out_path:
+        with open(out_path, "wb") as f:
+            f.write(tflite_quant_model)
+
+    return tflite_quant_model
+
+
+def representative_data(data, image_shape, N, channels):
+    """Obtains dataset from data, samples N samples, and returns those
+samples reshaped to image_shape.
+
+    """
+    path = data.get('test', data['val'])
+    f = []
+    for p in path if isinstance(path, list) else [path]:
+        if isdir(p):
+            f += glob(join(p, "**", "*.*"), recursive=True)
+        else:
+            f += [t if isabs(t) else join(dirname(p), t)
+                  for t in open(p).read().splitlines()]
+    shuffle(f)
+    for fpth in f[:N]:
+        im = imread(fpth)
+        if im.shape[0] != image_shape[0] or im.shape[1] != image_shape[1]:
+            im = resize(im, image_shape[::-1])
+        if im.shape[-1] != channels:
+            assert channels == 1
+            im = im.mean(-1, keepdims=True)
+        yield [im[None].astype(float32) / 255]
+
+
+def phony_data(image_shape, channels):
+    for _ in range(2):
+        yield [rand(1, *image_shape, channels).astype(float32)]
+
+
+def main(args=None):
+    return run(**vars(parse_opt(args)))

synet/sabre.py

@@ -0,0 +1,7 @@
+"""katana.py includes imports which are compatible with Katana, and
+layer definitions that are only compatible with Katana.  However,
+Katana's capabilities are currently a subset of all other chip's
+capabilities, so it includes only imports for now."""
+
+from .layers import Conv2dInvertedResidual, Head
+from .base import askeras, Conv2d, Cat, ReLU, BatchNorm, Grayscale

synet/tflite_utils.py

@@ -0,0 +1,349 @@
+#!/usr/bin/env python
+"""This module exists to hold all tflite related processing.  The main
+benefit of keeping this in a seperate modules is so that large
+dependencies (like ultralytics) need not be imported when simulating
+tflite execution (like for demos).  However, visualization
+(interpretation} of the model is left to ultralytics.  This module
+also serves as a reference for C and/or other implementations;
+however, do read any "Notes" sections in any function docstrings
+
+"""
+
+from argparse import ArgumentParser
+from typing import Optional, List
+
+from cv2 import (imread, rectangle, addWeighted, imwrite, resize,
+                 circle, putText, FONT_HERSHEY_TRIPLEX)
+from numpy import (newaxis, ndarray, int8, float32 as npfloat32,
+                   concatenate as cat, max as npmax, argmax, empty,
+                   array)
+from tensorflow import lite
+from torch import tensor, float32 as torchfloat32, sigmoid, tensordot, \
+    repeat_interleave
+from torchvision.ops import nms
+
+
+def parse_opt(args=None):
+    parser = ArgumentParser()
+    parser.add_argument('tflite')
+    parser.add_argument('img')
+    parser.add_argument('--conf-thresh', type=float, default=.25)
+    parser.add_argument('--iou-thresh', type=float, default=.5)
+    parser.add_argument('--backend', default='ultralytics')
+    parser.add_argument('--task', default='segment')
+    parser.add_argument('--image-shape', nargs=2, type=int, default=None)
+    return parser.parse_args(args=args)
+
+
+def tf_run(tflite, img, conf_thresh=.25, iou_thresh=.7,
+           backend='ultralytics', task='segment', image_shape=None):
+    """Run a tflite model on an image, including post-processing.
+
+    Loads the tflite, loads the image, preprocesses the image,
+    evaluates the tflite on the pre-processed image, and performs
+
+    post-processing on the tflite output with a given confidence and
+    iou threshold.
+
+    Parameters
+    ----------
+    tflite : str or buffer
+        Path to tflite file, or a raw tflite buffer
+    img : str or ndarray
+        Path to image to evaluate on, or the image as read by cv2.imread.
+    conf_thresh : float
+        Confidence threshold applied before NMS
+    iou_thresh : float
+        IoU threshold for NMS
+    backend : {"ultralytics"}
+        The backend which is used.  For now, only "ultralytics" is supported.
+    task : {"classify", "detect", "segment", "pose"}
+        The computer vision task to which the tflite model corresponds.
+
+    Returns
+    -------
+    ndarray or tuple of ndarrys
+        Return the result of running preprocessing, tflite evaluation,
+        and postprocessing on the input image.  Segmentation models
+        produce two outputs as a tuple.
+
+    """
+
+    # initialize tflite interpreter.
+    interpreter = lite.Interpreter(
+        **{"model_path" if isinstance(tflite, str)
+           else "model_content": tflite,
+           'experimental_op_resolver_type':
+           lite.experimental.OpResolverType.BUILTIN_REF}
+    )
+
+    # read the image if given as path
+    if isinstance(img, str):
+        img = imread(img)
+
+    if image_shape is not None:
+        img = resize(img, image_shape[::-1])
+
+    # make image RGB (not BGR) channel order, BCHW dimensions, and
+    # in the range [0, 1].  cv2's imread reads in BGR channel
+    # order, with dimensions in Height, Width, Channel order.
+    # Also, imread keeps images as integers in [0,255].  Normalize
+    # to floats in [0, 1].  Also, model expects a batch dimension,
+    # so add a dimension at the beginning
+    img = img[newaxis, ..., ::-1] / 255
+    # FW TEAM NOTE: It might be strange converting to float here, but
+    # the model might have been quantized to use a subset of the [0,1]
+    # range, i.e. 220 could map to 255
+
+    # Run tflite interpreter on the input image
+    preds = run_interpreter(interpreter, img)
+
+    if task == 'classify':
+        return preds
+
+    # Procces the tflite output to be one tensor
+    preds = concat_reshape(preds, task)
+
+    # perform nms
+    return apply_nms(preds, conf_thresh, iou_thresh)
+
+
+def run_interpreter(interpreter: Optional[lite.Interpreter],
+                    input_arr: ndarray) -> List[ndarray]:
+    """Evaluating tflite interpreter on input data
+
+    Parameters
+    ----------
+    interpreter : Interpreter
+        the tflite interpreter to run
+    input_arr : 4d ndarray
+        tflite model input with shape (batch, height, width, channels)
+
+    Returns
+    -------
+    list
+        List of output arrays from running interpreter.  The order and
+        content of the output is specific to the task and if model
+        outputs xywh or xyxy.
+    """
+
+    interpreter.allocate_tensors()
+    in_scale, in_zero = interpreter.get_input_details()[0]['quantization']
+    out_scale_zero_index = [(*detail['quantization'], detail['index'])
+                            for detail in
+                            sorted(interpreter.get_output_details(),
+                                   key=lambda x:x['name'])]
+    # run tflite on image
+    assert interpreter.get_input_details()[0]['index'] == 0
+    assert interpreter.get_input_details()[0]['dtype'] is int8
+    interpreter.set_tensor(0, (input_arr / in_scale + in_zero).astype(int8))
+    interpreter.invoke()
+    # indexing below with [0] removes the batch dimension, which is always 1.
+    return [(interpreter.get_tensor(index)[0].astype(npfloat32) - zero) * scale
+            for scale, zero, index in out_scale_zero_index]
+
+
+def concat_reshape(model_output: List[ndarray],
+                   task: str,
+                   xywh: Optional[bool] = False,
+                   classes_to_index: Optional[bool] = True
+                   ) -> ndarray:
+    """Concatenate, reshape, and transpose model output to match pytorch.
+
+    This method reordering the tflite output structure to be fit to run
+    post process such as NMS etc.
+
+    Parameters
+    ----------
+    model_output : list
+        Output from running tflite.
+    task : {"classify", "detect", "segment", "pose"}
+        The task the model performs.
+    xywh : bool, default=False
+        If true, model output should be converted to xywh.  Only use for
+        python evaluation.
+    classes_to_index : bool, default=True
+        If true, convert the classes output logits to single class index
+
+    Returns
+    -------
+    ndarray or list
+        Final output after concatenating and reshaping input.  Returns
+        an ndarray for every task except "segment" which returns a
+        tupule of two arrays.
+
+    Notes
+    -----
+    The python implementation here concats all output before applying
+    nms.  This is to mirror the original pytorch implementation.  For
+    a more efficient implementation, you may want to perform
+    confidence thresholding and nms on the boxes and scores, masking
+    other tensor appropriately, before reshaping and concatenating.
+
+    """
+
+    # interperate input tuple of tensors based on task.  Though
+    # produced tflite always have output names like
+    # "StatefulPartitionedCall:0", the numbers at the end are infact
+    # alphabetically ordered by the final layer name for each output,
+    # even though those names are discarded.  Hence, the following
+    # variables are nammed to match the corresponding output layer
+    # name and always appear in alphabetical order.
+    if task == "pose":
+        box1, box2, cls, kpts, pres = model_output
+        _, num_kpts, _ = kpts.shape
+    if task == "segment":
+        box1, box2, cls, proto, seg = model_output
+    if task == "detect":
+        box1, box2, cls = model_output
+
+    # obtain class confidences.
+    if classes_to_index:
+        # for yolov5, treat objectness seperately
+        cls = (npmax(cls, axis=1, keepdims=True),
+               argmax(cls, axis=1, keepdims=True))
+    else:
+        cls = cls,
+
+    # xywh is only necessary for python evaluation.
+    if xywh:
+        bbox_xy_center = (box1 + box2) / 2
+        bbox_wh = box2 - box1
+        bbox = cat([bbox_xy_center, bbox_wh], -1)
+    else:
+        bbox = cat([box1, box2], -1)
+
+    # return final concatenated output
+    # FW TEAM NOTE: Though this procedure creates output consistent
+    # with the original pytorch behavior of these models, you probably
+    # want to do something more clever, i.e. perform NMS reading from
+    # the arrays without concatenating.  At the very least, maybe do a
+    # confidence filter before trying to copy the full tensors.  Also,
+    # future models might have several times the output size, so keep
+    # that in mind.
+    if task == "segment":
+        # FW TEAM NOTE: the second element here move channel axis to
+        # beginning in line with pytorch behavior.  Maybe not relevent.
+
+        # FW TEAM NOTE: the proto array is HUGE (HxWx64).  You
+        # probably want to compute individual instance masks for your
+        # implementation.  See the YOLACT paper on arxiv.org:
+        # https://arxiv.org/abs/1904.02689.  Basically, for each
+        # instance that survives NMS, generate the segmentation (only
+        # HxW for each instance) by taking the iner product of seg
+        # with each pixel in proto.
+        return cat((bbox, *cls, seg), axis=-1), proto
+    if task == 'pose':
+        return cat((bbox, *cls, cat((kpts, pres), -1
+                                    ).reshape(-1, num_kpts * 3)),
+                   axis=-1)
+    if task == 'detect':
+        return cat((bbox, *cls), axis=-1)
+
+
+def apply_nms(preds: ndarray, conf_thresh: float, iou_thresh: float):
+    """Apply NMS on ndarray prepared model output
+
+    preds : ndarray or tuple of ndarray
+        prepared model output.  Is a tuple of two arrays for "segment" task
+    conf_thresh : float
+        confidence threshold applied before NMS.
+
+    Returns
+    -------
+    ndarray or tuple of ndarray
+        same structure as preds, but with some values suppressed (removed).
+
+    Notes
+    -----
+    This function converts ndarrays to pytorch tensors for two reasons:
+     - the nms code requires torch tensor inputs
+     - the output format becomes identical to that used by
+       ultralytics, and so can be passed to an ultralytics visualizer.
+
+    Also, as mentioned in the concat_reshape function, you may want to
+    perform nms and thresholding before combining all the output.
+
+    """
+
+    # THIS FUNCTION IS CURRENTLY HARD-CODED FOR SINGLE CLASS
+
+    # segmentation task returns a tuple of (preds, proto)
+    if isinstance(preds, tuple):
+        is_tuple = True
+        preds, proto = preds
+    else:
+        is_tuple = False
+
+    # perform confidence thresholding, and convert to tensor for nms.
+    # The trickiness here is that yolov5 has an objectness score plus
+    # per-class probabilities while ultralytics has just per-class
+    # scores.  Yolov5 uses objectness for confidence thresholding, but
+    # then uses objectness * per-class probablities for confidences
+    # therafter.
+    preds = tensor(preds[preds[:, 4] > conf_thresh], dtype=torchfloat32)
+
+    # Perform NMS
+    # https://pytorch.org/vision/stable/generated/torchvision.ops.nms.html
+    preds = preds[nms(preds[:, :4], preds[:, 4], iou_thresh)]
+    return (preds, tensor(proto)) if is_tuple else preds
+
+
+def build_masks(preds, proto):
+    # contract mask embeddings with proto
+    #                      (  N x k   dot   k x h x w  )
+    masks = sigmoid(tensordot(preds[:, 6:], proto, dims=1))
+    # upsamle mask
+    for dim in 1, 2:
+        masks = repeat_interleave(masks, repeats=8, dim=dim)
+    # clip mask to box
+    for (x1, y1, x2, y2), mask in zip(preds[:, :4], masks):
+        # integer math may be off-by-one near boarder in this
+        # application.
+        mask[:int(y1)] = 0
+        mask[int(y2):] = 0
+        mask[:, :int(x1)] = 0
+        mask[:, int(x2):] = 0
+    return preds[:, :6], masks
+
+
+def main(args=None):
+    # read options, run model
+    opt = parse_opt(args)
+    img = opt.img = imread(opt.img)
+    if opt.image_shape is not None:
+        img = opt.img = resize(opt.img, opt.image_shape[::-1])
+        opt.image_shape = None
+
+    preds = tf_run(**vars(opt))
+    if opt.task == 'segment':
+        preds, masks = build_masks(*preds)
+
+        # shrink mask if upsamle gave too large of area
+        for imdim, maskdim in (0, 1), (1, 2):
+            extra, carry = divmod(masks.shape[maskdim] - img.shape[imdim], 2)
+            if extra == carry == 0:
+                continue
+            masks = masks[(*(slice(None),)*maskdim, slice(extra, -(extra+carry)))]
+
+        # visualize masks and rectangles
+        img_overlay = img.copy()
+        img_overlay[masks.max(0).values > .5] = (0, 255, 0)
+        img = addWeighted(img, .5, img_overlay, .5, 0)
+    elif opt.task == 'pose':
+        for x1, y1, x2, y2, conf, cls, *kpts in preds:
+            for x, y, p in zip(kpts[0::3], kpts[1::3], kpts[2::3]):
+                if p > .5:
+                    circle(img, (int(x), int(y)), 3, (255, 0, 0), -1)
+    elif opt.task != 'classify':
+        for x1, y1, x2, y2, *cls in preds:
+            rectangle(img, (int(x1), int(y1)),
+                      (int(x2), int(y2)),
+                      (0, 0, 255), 2)
+    elif opt.task == 'classify':
+        putText(img, str(*preds), (20, 40), FONT_HERSHEY_TRIPLEX, 1.0, (0, 0, 0))
+    imwrite(opt.task+'.png', img)
+
+
+if __name__ == '__main__':
+    main()

synet/ultralytics_patches.py

@@ -0,0 +1,3 @@
+from .backends.ultralytics import *
+print("WARNING: ultralytics_patches.py exists for backwards model "
+      "compatibility only.  Do not import this module if possible.")

synet/zoo/__init__.py

@@ -0,0 +1,35 @@
+from os import listdir
+from os.path import abspath, dirname, join, isfile, commonpath
+from urllib import request
+
+
+WEIGHT_URL_ROOT = "http://profiler/"
+
+
+def in_zoo(model, backend):
+    """Return True if model refers to something in the SyNet zoo."""
+    # check if absolute path to model in the zoo was given
+    if isfile(model):
+        return dirname(__file__) == commonpath((__file__, abspath(model)))
+    # otherwise check if name is relative to zoo dir.
+    return isfile(join(dirname(__file__), backend, model))
+
+
+def get_config(model, backend):
+    """Return the path to a model.  Check the zoo if necessary."""
+    if isfile(model):
+        return model
+    return join(dirname(__file__), backend, model)
+
+
+def get_weights(model, backend):
+    if isfile(model):
+        return model
+    with request.urlopen(join(WEIGHT_URL_ROOT, backend, model)) as remotefile:
+        with open(model, 'wb') as localfile:
+            localfile.write(remotefile.read())
+    return model
+
+
+def get_configs(backend):
+    return listdir(join(dirname(__file__), backend))

synet/zoo/ultralytics/sabre-detect-vga.yaml

@@ -0,0 +1,29 @@
+nc: 1  # number of classes
+#kpt_shape: [17, 3]
+depth_multiple: 1  # model depth multiple
+width_multiple: 1  # layer channel multiple
+chip: sabre
+image_shape: [480, 640]
+# anchors:
+#     # autogenerated by yolo
+#   - [0,0, 0,0, 0,0] # P3/8
+#   - [0,0, 0,0, 0,0] # P4/16
+#   - [0,0, 0,0, 0,0] # P5/32
+backbone:
+  # [from, number, module, args]
+  #src num layer               params                  id  rf  notes
+  [[-1, 1, InvertedResidual, [3, 4, 12, 2]],   #  0 c1   1 stride -> 2
+   [-1, 1, InvertedResidual, [12, 4, 48, 2]],     #  1 c2   3 stride -> 4
+   [-1, 1, InvertedResidual, [48, 4, 48, 2]],    #  2    7 stride -> 8
+   [-1, 2, InvertedResidual, [48, 6, 48]],       #  3 c3
+   [-1, 1, InvertedResidual, [48, 5, 64, 2]],    #  4     15 stride -> 16
+   [-1, 2, InvertedResidual, [64, 4, 64]],           #  5 c4  47
+   [-1, 1, InvertedResidual, [64, 3, 64, 2]],    #  6     63 stride -> 32
+   [-1, 2, InvertedResidual, [64, 2]]]           #  7 c5 127
+
+# YOLOv5 v6.0 head
+head:
+  [[ 5, 1, Head,             [64, 64, 3]],       #  8 o4  95
+   [ 7, 1, Head,             [64, 64, 3]],       #  9 o5 191
+   [[ 8, 9], 1, Detect, [nc, [64, 64], 2]]   # 43 Detect(P4-P6)
+   ]                              # rfs 127 255

synet/zoo/ultralytics/sabre-keypoint-vga.yaml

@@ -0,0 +1,29 @@
+nc: 1  # number of classes
+kpt_shape: [17, 3]
+depth_multiple: 1  # model depth multiple
+width_multiple: 1  # layer channel multiple
+chip: sabre
+image_shape: [480, 640]
+# anchors:
+#     # autogenerated by yolo
+#   - [0,0, 0,0, 0,0] # P3/8
+#   - [0,0, 0,0, 0,0] # P4/16
+#   - [0,0, 0,0, 0,0] # P5/32
+backbone:
+  # [from, number, module, args]
+  #src num layer               params                  id  rf  notes
+  [[-1, 1, InvertedResidual, [3, 4, 12, 2]],   #  0 c1   1 stride -> 2
+   [-1, 1, InvertedResidual, [12, 4, 48, 2]],     #  1 c2   3 stride -> 4
+   [-1, 1, InvertedResidual, [48, 4, 48, 2]],    #  2    7 stride -> 8
+   [-1, 2, InvertedResidual, [48, 5, 48]],       #  3 c3
+   [-1, 1, InvertedResidual, [48, 4, 64, 2]],    #  4     15 stride -> 16
+   [-1, 2, InvertedResidual, [64, 3, 64]],           #  5 c4  47
+   [-1, 1, InvertedResidual, [64, 2, 64, 2]],    #  6     63 stride -> 32
+   [-1, 2, InvertedResidual, [64, 1]]]           #  7 c5 127
+
+# YOLOv5 v6.0 head
+head:
+  [[ 5, 1, Head,             [64, 64, 3]],       #  8 o4  95
+   [ 7, 1, Head,             [64, 64, 3]],       #  9 o5 191
+   [[ 8, 9], 1, Pose, [nc, [17,3], [64, 64], 2]]   # 43 Detect(P4-P6)
+   ]                              # rfs 127 255

synet/zoo/ultralytics/sabre-segment-vga.yaml

@@ -0,0 +1,29 @@
+nc: 1  # number of classes
+#kpt_shape: [17, 3]
+depth_multiple: 1  # model depth multiple
+width_multiple: 1  # layer channel multiple
+chip: sabre
+image_shape: [480, 640]
+# anchors:
+#     # autogenerated by yolo
+#   - [0,0, 0,0, 0,0] # P3/8
+#   - [0,0, 0,0, 0,0] # P4/16
+#   - [0,0, 0,0, 0,0] # P5/32
+backbone:
+  # [from, number, module, args]
+  #src num layer               params                  id  rf  notes
+  [[-1, 1, InvertedResidual, [3, 4, 12, 2]],   #  0 c1   1 stride -> 2
+   [-1, 1, InvertedResidual, [12, 4, 48, 2]],     #  1 c2   3 stride -> 4
+   [-1, 1, InvertedResidual, [48, 4, 48, 2]],    #  2    7 stride -> 8
+   [-1, 2, InvertedResidual, [48, 5, 48]],       #  3 c3
+   [-1, 1, InvertedResidual, [48, 4, 64, 2]],    #  4     15 stride -> 16
+   [-1, 2, InvertedResidual, [64, 3, 64]],           #  5 c4  47
+   [-1, 1, InvertedResidual, [64, 2, 64, 2]],    #  6     63 stride -> 32
+   [-1, 2, InvertedResidual, [64, 1]]]           #  7 c5 127
+
+# YOLOv5 v6.0 head
+head:
+  [[ 5, 1, Head,             [64, 64, 3]],       #  8 o4  95
+   [ 7, 1, Head,             [64, 64, 3]],       #  9 o5 191
+   [[ 8, 9], 1, Segment, [nc, 32, 96, [64, 64], 2]]   # 43 Detect(P4-P6)
+   ]                              # rfs 127 255

tests/test_demosaic.py

@@ -0,0 +1,17 @@
+from torch import rand
+from torch.nn.init import uniform_
+
+from synet.base import Conv2d
+from synet.demosaic import reshape_conv
+
+
+def test_reshape_conv():
+    conv = Conv2d(3, 13, 3, 2)
+    for param in conv.parameters():
+        uniform_(param, -1)
+    reshaped_conv = reshape_conv(conv)
+    inp = rand(3, 480, 640)
+    reshaped_inp = inp.reshape(3, 240, 2, 320, 2
+                               ).permute(2, 4, 0, 1, 3
+                                         ).reshape(12, 240, 320)
+    assert (reshaped_conv(reshaped_inp) - conv(inp)).abs().max() < 1e-5

tests/test_keras.py

@@ -0,0 +1,210 @@
+from numpy import absolute
+from torch import rand
+from torch.nn.init import uniform_
+
+from synet.base import askeras
+
+
+BATCH_SIZE = 2
+IN_CHANNELS = 5
+OUT_CHANNELS = 7
+SHAPES = [(i, i) for i in range(4, 8)]
+MAX_DIFF = -1
+TOLERANCE = 2e-4
+
+
+def diff_arr(out1, out2):
+    """compare two arrays.  Return the max difference."""
+    if isinstance(out1, (list, tuple)):
+        assert isinstance(out2, (list, tuple))
+        return max(diff_arr(o1, o2) for o1, o2 in zip(out1, out2))
+    assert all(s1 == s2 for s1, s2 in zip(out1.shape, out2.shape)), \
+        (out1.shape, out2.shape)
+    return absolute(out1 - out2).max()
+
+
+def t_actv_to_k(actv):
+    if isinstance(actv, (tuple, list)):
+        return [t_actv_to_k(a) for a in actv]
+    if len(actv.shape) == 4:
+        tp = 0, 2, 3, 1
+    elif len(actv.shape) == 3:
+        tp = 0, 2, 1
+    elif len(actv.shape) == 2:
+        tp = 0, 1
+    return actv.detach().numpy().transpose(*tp)
+
+
+def k_to_numpy(actv):
+    if isinstance(actv, (list, tuple)):
+        return [k_to_numpy(k) for k in actv]
+    if hasattr(actv, "numpy"):
+        return actv.numpy()
+    return actv
+
+
+def validate_layer(layer, torch_inp, **akwds):
+    """Given synet layer, test on some torch input activations and
+return max error between two output activations
+
+    """
+    tout = layer(torch_inp[:])
+    with askeras(imgsz=torch_inp[0].shape[-2:], **akwds):
+        kout = k_to_numpy(layer(t_actv_to_k(torch_inp)))
+    if isinstance(tout, dict):
+        assert len(tout) == len(kout)
+        return max(diff_arr(t_actv_to_k(tout[key]), kout[key])
+                   for key in tout)
+    elif isinstance(tout, list):
+        assert len(tout) == len(kout)
+        return max(diff_arr(t_actv_to_k(t), k)
+                   for t, k in zip(tout, kout))
+    return diff_arr(t_actv_to_k(tout), kout)
+
+
+def validate(layer, batch_size=BATCH_SIZE,
+             in_channels=IN_CHANNELS, shapes=SHAPES, **akwds):
+    """Run validate_layer on a set of random input shapes.  Prints the max
+difference between all configurations.
+
+    """
+    for param in layer.parameters():
+        uniform_(param, -1)
+    max_diff = max(validate_layer(layer,
+                                  [rand(batch_size, in_channels, *s)*2-1
+                                   for s in shape]
+                                  if len(shape) and isinstance(shape[0], tuple)
+                                  else rand(batch_size, in_channels, *shape
+                                            )*2-1,
+                                  **akwds)
+                   for shape in shapes)
+    print("max_diff:", max_diff)
+    assert max_diff < TOLERANCE
+
+
+def test_conv2d():
+    from synet.base import Conv2d
+    print("testing Conv2d")
+    in_channels = 12
+    out_channels = 24
+    for bias in True, False:
+        for kernel, stride in ((1, 1), (2, 1), (3, 1), (3, 2), (4, 1),
+                               (4, 2), (4, 3)):
+            for padding in True, False:
+                for groups in 1, 2, 3:
+                    validate(Conv2d(in_channels, out_channels, kernel,
+                                    stride, bias, padding),
+                             in_channels=in_channels)
+
+def test_dw_conv2d():
+    from synet.layers import DepthwiseConv2d
+    print("testing dw Conv2d")
+    channels = 32
+    for bias in True, False:
+        for kernel, stride in ((1, 1), (2, 1), (3, 1), (3, 2), (4, 1),
+                               (4, 2), (4, 3)):
+            for padding in True, False:
+                validate(DepthwiseConv2d(channels, kernel,
+                                stride, bias, padding),
+                            in_channels=channels)
+
+
+def test_convtranspose():
+    from synet.base import ConvTranspose2d
+    validate(ConvTranspose2d(IN_CHANNELS, OUT_CHANNELS, 2, 2, 0, bias=True))
+
+
+def test_relu():
+    from synet.base import ReLU
+    validate(ReLU(.6))
+
+
+def test_upsample():
+    from synet.base import Upsample
+    for scale_factor in 1, 2, 3:
+        for mode in Upsample.allowed_modes:
+            validate(Upsample(scale_factor, mode))
+
+
+def test_globavgpool():
+    from synet.base import GlobalAvgPool
+    validate(GlobalAvgPool())
+
+
+def test_dropout():
+    from synet.base import Dropout
+    for p in 0.0, 0.5, 1.0:
+        for inplace in True, False:
+            layer = Dropout(p, inplace=inplace)
+            layer.eval()
+            validate(layer)
+
+
+def test_linear():
+    from synet.base import Linear
+    for bias in True, False:
+        validate(Linear(IN_CHANNELS, OUT_CHANNELS, bias), shapes=[()])
+
+
+def test_batchnorm():
+    from synet.base import BatchNorm
+    validate(BatchNorm(IN_CHANNELS), train=True)
+
+
+def test_ultralytics_detect():
+    from synet.backends.ultralytics import Detect
+    for sm_split in ((True, None), (2, True)):
+        layer = Detect(80, (IN_CHANNELS, IN_CHANNELS), *sm_split)
+        layer.eval()
+        layer.export = True
+        layer.format = "tflite"
+        layer.stride[0], layer.stride[1] = 1, 2
+        validate(layer,
+                 shapes=[((4, 6), (2, 3)),
+                         ((5, 7), (3, 4)),
+                         ((6, 8), (3, 4))],
+                 xywh=True)
+
+
+def test_ultralytics_pose():
+    from synet.backends.ultralytics import Pose
+    for sm_split in ((True, None), (2, True)):
+        for kpt_shape in ([17, 2], [17, 3]):
+            layer = Pose(80, kpt_shape, (IN_CHANNELS, IN_CHANNELS), *sm_split)
+            layer.eval()
+            layer.export = True
+            layer.format = "tflite"
+            layer.stride[0], layer.stride[1] = 1, 2
+            validate(layer,
+                     shapes=[((4, 6), (2, 3)),
+                             ((5, 7), (3, 4)),
+                             ((6, 8), (3, 4))],
+                     xywh=True)
+
+
+def test_ultralytics_segment():
+    from synet.backends.ultralytics import Segment
+    layer = Segment(nc=80, nm=32, npr=256, ch=(IN_CHANNELS, IN_CHANNELS))
+    layer.eval()
+    layer.export = True
+    layer.format = "tflite"
+    layer.stride[0], layer.stride[1] = 1, 2
+    validate(layer,
+             shapes=[(( 4,  4), (2, 2)),
+                     (( 8,  8), (4, 4)),
+                     ((12, 12), (6, 6))],
+             xywh=True)
+
+
+def test_ultralytics_classify():
+    from synet.backends.ultralytics import Classify
+    layer = Classify(None, c1=IN_CHANNELS, c2=OUT_CHANNELS)
+    layer.eval()
+    layer.export = True
+    layer.format = 'tflite'
+    validate(layer)
+
+
+def test_channelslice():
+    from synet.base import ChannelSlice
+    validate(ChannelSlice(slice(4, 8)), in_channels=12)

tests/test_ultralytics.py

@@ -0,0 +1,18 @@
+from os import chdir
+from synet.backends import get_backend
+from synet.quantize import main
+
+
+backend = get_backend("ultralytics")
+
+
+def test_quantize(tmp_path):
+    chdir(tmp_path)
+    for config in backend.get_configs():
+        main(("--backend=ultralytics",
+              "--model="+config,
+              "--number=1"))
+    main(("--backend=ultralytics",
+          "--model="+config,
+          "--number=1",
+          "--image-shape", "321", "319"))