If you pattern match on the shape (and the type for good measure) of the %StbImage{…} your def pipeline(…) could read as:
def pipeline(paths, batch_size, target_height, target_width) do
paths
|> Enum.shuffle()
|> Task.async_stream(&parse_image/1)
|> Stream.filter(fn
# {:ok, {%StbImage{}, _}} -> true
{:ok, {%StbImage{data: _, shape: {_, _, 3}, type: {:u, 8}}, _}} -> true
_ -> false
end)
|> Stream.map(&to_tensors(&1, target_height, target_width))
|> Stream.chunk_every(batch_size, batch_size, :discard)
|> Stream.map(fn chunks ->
{img_chunk, label_chunk} = Enum.unzip(chunks)
{Nx.stack(img_chunk), Nx.stack(label_chunk)}
end)
end
The commented line is just to show what I changed. Cheers!