View Source Bumblebee.Vision (Bumblebee v0.2.0)

High-level tasks related to vision.

Link to this section Summary


Builds serving for image classification.

Link to this section Types

@type image() :: Nx.Container.t()

A term representing an image.

Either Nx.Tensor in HWC order or a struct implementing Nx.Container and resolving to such tensor.

Link to this type


View Source
@type image_classification_input() :: image()
Link to this type


View Source
@type image_classification_output() :: %{
  predictions: [image_classification_prediction()]
Link to this type


View Source
@type image_classification_prediction() :: %{score: number(), label: String.t()}

Link to this section Functions

Link to this function

image_classification(model_info, featurizer, opts \\ [])

View Source
@spec image_classification(
) :: Nx.Serving.t()

Builds serving for image classification.

The serving accepts image_classification_input/0 and returns image_classification_output/0. A list of inputs is also supported.



  • :top_k - the number of top predictions to include in the output. If the configured value is higher than the number of labels, all labels are returned. Defaults to 5

  • :compile - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys:

    • :batch_size - the maximum batch size of the input. Inputs are optionally padded to always match this batch size

    It is advised to set this option in production and also configure a defn compiler using :defn_options to maximally reduce inference time.

  • :defn_options - the options for JIT compilation. Defaults to []



{:ok, resnet} = Bumblebee.load_model({:hf, "microsoft/resnet-50"})
{:ok, featurizer} = Bumblebee.load_featurizer({:hf, "microsoft/resnet-50"})

serving = Bumblebee.Vision.image_classification(resnet, featurizer)

image = ..., image)
#=> %{
#=>   predictions: [
#=>     %{label: "Egyptian cat", score: 0.979233980178833},
#=>     %{label: "tabby, tabby cat", score: 0.00679466687142849},
#=>     %{label: "tiger cat", score: 0.005290505941957235},
#=>     %{label: "lynx, catamount", score: 0.004550771787762642},
#=>     %{label: "Siamese cat, Siamese", score: 1.1611092486418784e-4}
#=>   ]
#=> }