defmodule Evision.Zoo.ImageSegmentation.PPHumanSeg do
@moduledoc """
PP-Human Segmentation model.
"""
@doc """
Default configuration.
"""
@spec default_config :: map()
def default_config do
%{
backend: Evision.Constant.cv_DNN_BACKEND_OPENCV(),
target: Evision.Constant.cv_DNN_TARGET_CPU(),
}
end
@doc """
Customizable parameters from smart cell.
"""
@spec smartcell_params() :: Evision.Zoo.smartcell_params()
def smartcell_params() do
[]
end
@doc """
Initialize model.
##### Positional arguments
- **model**: `String.t()` | `:default_model` | `:quant_model`.
- When `model` is a string, it will be treat as the path to a weight file
and `init/2` will load the model from it.
- When `model` is either `:default_model` or `:quant_model`, `init/2` will
download and load the predefined model.
##### Keyword arguments
- **cache_dir**: `String.t()`.
Path to the cache directory.
Optional. Defaults to `:filename.basedir(:user_cache, "", ...)`
- **backend**: `integer()`.
Specify the backend.
Optional. Defaults to `Evision.Constant.cv_DNN_BACKEND_OPENCV()`.
- **target**: `integer()`.
Specify the target.
Optional. Defaults to `Evision.Constant.cv_DNN_TARGET_CPU()`.
"""
@spec init(binary | :default_model | :quant_model, nil | Keyword.t()) :: {:error, String.t()} | Evision.DNN.Net.t()
def init(model, opts \\ [])
def init(model_type, opts) when model_type in [:default_model, :quant_model] do
{model_url, filename} = model_info(model_type)
cache_dir = opts[:cache_dir]
with {:ok, local_path} <- Evision.Zoo.download(model_url, filename, cache_dir: cache_dir) do
init(local_path, opts)
else
{:error, msg} ->
raise msg
end
end
def init(model_path, opts) when is_binary(model_path) do
config = default_config()
backend = opts[:backend] || config[:backend]
target = opts[:target] || config[:target]
net = Evision.DNN.readNet(model_path)
Evision.DNN.Net.setPreferableBackend(net, backend)
Evision.DNN.Net.setPreferableTarget(net, target)
net
end
@doc """
Inference.
##### Positional arguments
- **self**: `Evision.DNN.Net.t()`.
An initialized PPHumanSeg model.
- **image**: `Evision.Mat.maybe_mat_in()`.
Input image.
"""
@spec infer(Evision.DNN.Net.t(), Evision.Mat.maybe_mat_in()) :: Evision.Mat.t()
def infer(self=%Evision.DNN.Net{}, image) do
inputBlob = preprocess(image)
Evision.DNN.Net.setInput(self, inputBlob)
outputBlob = Evision.DNN.Net.forward(self, outputName: "save_infer_model/scale_0.tmp_1")
outputBlob =
if is_list(outputBlob) do
[outputBlob] = outputBlob
outputBlob
else
outputBlob
end
# todo: use Evision.Backend when Nx.slice is implemented
result = Evision.Mat.to_nx(Evision.Mat.squeeze(outputBlob), Nx.BinaryBackend)
Evision.Mat.from_nx(Nx.as_type(Nx.argmax(result, axis: 0), :u8))
end
@doc """
Preprocessing the input image.
`infer/2` will call this function automatically.
##### Positional arguments
- **image**: `Evision.Mat.maybe_mat_in()`.
Input image.
"""
@spec preprocess(Evision.Mat.maybe_mat_in()) :: Evision.Mat.t()
def preprocess(image) do
image
|> Evision.Mat.as_type(:f32)
|> Evision.resize({192, 192})
|> Evision.Mat.to_nx()
|> Nx.divide(Nx.broadcast(Nx.tensor(255.0, backend: Evision.Backend), {192, 192, 3}))
|> Nx.subtract(mean())
|> Nx.divide(Nx.broadcast(std(), {192, 192, 3}))
|> Evision.Mat.from_nx_2d()
|> Evision.DNN.blobFromImage()
end
defp mean do
Evision.Mat.to_nx(Evision.Mat.literal([[[0.5, 0.5, 0.5]]], :f32))
end
defp std do
Evision.Mat.to_nx(Evision.Mat.literal([[[0.5, 0.5, 0.5]]], :f32))
end
@doc """
Visualize the result.
##### Positional arguments
- **image**: `Evision.Mat.maybe_mat_in()`.
Original image.
- **results**: `Evision.Mat.maybe_mat_in()`.
Results given by `infer/2`.
##### Keyword arguments
- **weight**: `number()`.
A number in `[0.0, 1.0]`. Defaults to `0.6`.
Specify the weight of the original image. The weight of the segmentation visualization image
will be `1 - weight`.
##### Return
A list that contains two images (`Evision.Mat.t()`).
- The first one is the original image with the segmentation overlay.
- The second one is the segmentation image.
"""
@spec visualize(Evision.Mat.maybe_mat_in(), Evision.Mat.maybe_mat_in(), Keyword.t()) :: list(Evision.Mat.t())
def visualize(image, results, opts \\ []) do
weight = opts[:weight] || 0.6
color_map = Evision.Mat.literal(Evision.Zoo.ImageSegmentation.color_map(256), :u8)
c1 = Evision.lut(results, color_map[[:all, 0]])
c2 = Evision.lut(results, color_map[[:all, 1]])
c3 = Evision.lut(results, color_map[[:all, 2]])
segmentation = Evision.merge([c1, c2, c3])
[Evision.addWeighted(image, weight, segmentation, 1 - weight, 0), segmentation]
end
@doc """
Model URL and filename of predefined model.
"""
@spec model_info(:default_model | :quant_model) :: {String.t(), String.t()}
def model_info(:default_model) do
{
"https://github.com/opencv/opencv_zoo/blob/0d619617a8e9a389150d8c76e417451a19468150/models/human_segmentation_pphumanseg/human_segmentation_pphumanseg_2023mar.onnx?raw=true",
"human_segmentation_pphumanseg_2023mar.onnx"
}
end
def model_info(:quant_model) do
{
"https://github.com/opencv/opencv_zoo/blob/0d619617a8e9a389150d8c76e417451a19468150/models/human_segmentation_pphumanseg/human_segmentation_pphumanseg_2023mar_int8?raw=true",
"human_segmentation_pphumanseg_2023mar_int8.onnx"
}
end
@doc """
Docs in smart cell.
"""
@spec docs() :: String.t()
def docs do
@moduledoc
end
@doc """
Smart cell tasks.
A list of variants of the current model.
"""
@spec smartcell_tasks() :: Evision.Zoo.smartcell_tasks()
def smartcell_tasks do
[
%{
id: "pp_humanseg",
label: "PP-HumanSeg",
docs_url: "https://github.com/opencv/opencv_zoo/tree/master/models/human_segmentation_pphumanseg",
params: smartcell_params(),
docs: docs()
},
%{
id: "pp_humanseg_quant",
label: "PP-HumanSeg (quant)",
docs_url: "https://github.com/opencv/opencv_zoo/tree/master/models/human_segmentation_pphumanseg",
params: smartcell_params(),
docs: docs()
},
]
end
@doc """
Generate quoted code from smart cell attrs.
"""
@spec to_quoted(map()) :: list()
def to_quoted(attrs) do
{backend, target} = Evision.Zoo.to_quoted_backend_and_target(attrs)
opts = [
backend: backend,
target: target
]
model =
case attrs["variant_id"] do
"pp_humanseg_quant" ->
:quant_model
_ ->
:default_model
end
[
quote do
model = Evision.Zoo.ImageSegmentation.PPHumanSeg.init(unquote(model), unquote(opts))
end,
quote do
image_input = Kino.Input.image("Image")
form = Kino.Control.form([image: image_input], submit: "Run")
frame = Kino.Frame.new()
form
|> Kino.Control.stream()
|> Stream.filter(& &1.data.image)
|> Kino.listen(fn %{data: %{image: image}} ->
Kino.Frame.render(frame, Kino.Markdown.new("Running..."))
{height, width} = {image.height, image.width}
image =
image.file_ref
|> Kino.Input.file_path()
|> File.read!()
|> Evision.Mat.from_binary({:u, 8}, height, width, 3)
results = Evision.Zoo.ImageSegmentation.PPHumanSeg.infer(model, image)
results = Evision.resize(results, {width, height}, interpolation: Evision.Constant.cv_INTER_NEAREST())
image = Evision.cvtColor(image, Evision.Constant.cv_COLOR_RGB2BGR())
vis_imgs = Evision.Zoo.ImageSegmentation.PPHumanSeg.visualize(image, results)
vis_imgs = Enum.map(vis_imgs, &Kino.Image.new(Evision.imencode(".png", &1), :png))
Kino.Frame.render(frame, Kino.Layout.grid(vis_imgs, columns: 2))
end)
Kino.Layout.grid([form, frame], boxed: true, gap: 16)
end
]
end
end
