module.yml

name: super_resolution

parameters:
  # set the super resolution model scale (x2/3/4) and name (ninasr_b0/1)
  sr_scale: 3
  sr_model: ninasr_b0
  # needs to retrieve the model output in pyfunc element
  sr_attribute: sr_frame
  # pipeline processing frame parameters
  frame:
    width: 640
    height: 360
  batch_size: 1

pipeline:
  elements:
    - element: nvinfer@attribute_model
      name: ${parameters.sr_model}
      model:
        remote:
          url: s3://savant-data/models/super_resolution/${parameters.sr_model}/${parameters.sr_model}_x${parameters.sr_scale}.zip
          checksum_url: s3://savant-data/models/super_resolution/${parameters.sr_model}/${parameters.sr_model}_x${parameters.sr_scale}.md5
          parameters:
            endpoint: https://eu-central-1.linodeobjects.com
        format: onnx
        model_file: ${parameters.sr_model}_x${parameters.sr_scale}.onnx
        input:
          shape: [3, 360, 640]
          scale_factor: 0.003921569
          offsets: [0.0, 0.0, 0.0]
        output:
          layer_names: [output]
          converter:
            module: savant.converter.raw_output
            class_name: ModelCudaRawOutputConverter
          attributes:
            - name: ${parameters.sr_attribute}
              # just a way to save model output before place on frame, no need to output
              internal: true
    - element: pyfunc
      module: samples.super_resolution.overlay
      class_name: SROverlay
      kwargs:
        # codec parameters for result stream (auxiliary stream)
        codec_params:
          codec: hevc