Add segmentation.RandomMaskPixel operator

[jantonguirao]

Signed-off-by: Joaquin Anton janton@nvidia.com

Why we need this PR?

Pick one, remove the rest

• It adds new feature needed to be able to perform a random crop operation where foreground is always represented up to a minimum ratio

What happened in this PR?

Fill relevant points, put NA otherwise. Replace anything inside []

• What solution was applied:
  Added a new operator segmentation.RandomMaskPixel that selects a pixel either randomly or to point to any of the available foreground pixels, specified by an input mask and threhold/value arguments
• Affected modules and functionalities:
  No existing functionality is affected
• Key points relevant for the review:
  All
• Validation and testing:
  Python tests
• Documentation (including examples):
  Operator documentation in schema

JIRA TASK: [DALI-1723]

[JanuszL]

Suggested change

	any position in the input or any position of a foreground pixel in the input, based on an
	at any position in the input or any position of a foreground pixel in the input, based on an

[JanuszL]

I don't know if we should use foreground or just non zero value, because the meaning of the foreground if not defined.

[JanuszL]

I wonder if it makes any sense not to shift?

[jantonguirao]

I think both things can make sense. If cropping out of bounds is OK, the user can let shape=None and any pixel would be considered a valid center. If the shape is specified, we'll assure that the cropping window is within bounds.

[mzient]

We've already had github requests for out-of-bounds cropping. Perhaps we should make this configurable.

[jantonguirao]

We already have out-of-bounds cropping, and this is already configured (providing a shape or not).

[JanuszL]

Shouldn't you do anything with the return value as well?

[jantonguirao]

not really, I don't need to infer the batch size or the number of dimensions.

[JanuszL]

Can you tell how fast/slow it is for the typical inputs?
Do we assume that masks are uniform regarding the size across the batch or they can vary?
If can vary maybe we can differently split the work across the threads. Ask each thread to create a mask for only one plan in each volume.
Does it make any sense to have it on the GPU?

[JanuszL]

And how that works with 3D masks?

[jantonguirao]

I can benchmark it. I wouldn't assume uniform shape.
SearchableRLEMask works on flat indices. We are not doing anything with the dimensions.
Regarding the GPU, encoding the mask is a rather serial task. Having a mask per plane would complicate things, so I wouldn't go there unless we see there is a performance bottleneck here. On top of that, Slice operator doesn't take GPU anchor/shape.

[JanuszL]

Ok, let us benchmark first and see if there is anything to fight for.

[mzient]

I don't know what this w was supposed to mean, but it read as width (which apprently it isn't).

Suggested change

	int64_t w = crop_sh[d] >> 1;
	int64_t half = crop_sh[d] >> 1;

[mzient]

Suggested change

	boundary::idx_clamp(center.data[d], w, mask_sh[d] - (crop_sh[d] - w));
	clamp(center.data[d], w, mask_sh[d] - (crop_sh[d] - w));

If you're just clamping, you don't need any fancy boundary handling (ordinary clamp will do).
Also, boundary clamp is an exclusive clamp (to size-1) and we want to have inclusive clamp.
Example:
mask_sh[d] = 100
crop_sh[d] = 10
w = 5
maximum valid value is 95 (the window will start at 90 and end at 99)

[mzient]

Sample-indexed GetArgument is a relic from the obsolete SampleWorkspace. It's inefficient and should not be used.

Suggested change

	for (int sample_idx = 0; sample_idx < nsamples; sample_idx++) {
	foreground_[sample_idx] = spec_.template GetArgument<int>("foreground", &ws, sample_idx);
	}
	GetPerSampleArgument(foreground_, "foreground", spec_, ws, nsamples);

[mzient]

Suggested change

	value_.resize(nsamples);
	for (int sample_idx = 0; sample_idx < nsamples; sample_idx++) {
	value_[sample_idx] = spec_.template GetArgument<int>("value", &ws, sample_idx);
	}
	GetPerSampleArgument(value_, "value", spec_, ws, nsamples);

[mzient]

Suggested change

	for (int sample_idx = 0; sample_idx < nsamples; sample_idx++) {
	threshold_[sample_idx] = spec_.template GetArgument<float>("threshold", &ws, sample_idx);
	}
	GetPerSampleArgument(threshold_, "threshold", spec_, ws, nsamples);

[mzient]

Nitpick: resize won't do much more than a similar check, so:

Suggested change

	if (rle_.empty()) {
	rle_.resize(thread_pool.size());
	}
	assert(rle_.size() == static_cast<size_t>(thread_pool.size()));
	rle_.resize(thread_pool.size());

[mzient]

Don't do that. This will cause wrong results for signed integer input and fractional threshold (and also when the value wraps after casting).

Suggested change

	T threshold = static_cast<T>(threshold_[sample_idx]);
	float threshold = threshold_[sample_idx];

[mzient]

I have mixed feelings about that one - consider that:
input is uint8
value is 256
magically, you select 0 as foreground.
Actually:

if (static_cast<int>(value) != values_[sample_idx])
   // behave as if there was no foregound at all

[jantonguirao]

Done

[JanuszL]

Suggested change

	R"code(All pixels equal to this value are interpreted as foreground.
	R"code(All pixels equal to this value are interpreted as foreground.

[JanuszL]

Suggested change

	R"code(If different than 0, the pixel position is sampled uniformly from all foreground pixels.
	If 0, the pixel position is sampled uniformly from all available pixels.)code",
	R"code(If different than 0, the pixel position is sampled uniformly from all foreground pixels.
	If 0, the pixel position is sampled uniformly from all available pixels.)code",

[JanuszL]

Suggested change

	auto &rle_mask = rle_[thread_id];
	if (has_value_) {
	T value = static_cast<T>(value_[sample_idx]);
	// checking if the value is representable by T, otherwise we
	// just fall back to pick a random pixel.
	if (static_cast<int>(value) == value_[sample_idx]) {
	rle_mask.Init(
	mask, [value](const T &x) { return x == value; });
	if (rle_mask.count() > 0) {
	auto dist = std::uniform_int_distribution<int64_t>(0, rle_mask.count() - 1);
	flat_idx = rle_mask.find(dist(rng));
	}
	}
	} else {
	float threshold = threshold_[sample_idx];
	rle_mask.Init(
	mask, [threshold](const T &x) { return x > threshold; });
	if (rle_mask.count() > 0) {
	auto dist = std::uniform_int_distribution<int64_t>(0, rle_mask.count() - 1);
	flat_idx = rle_mask.find(dist(rng));
	}
	}
	auto &rle_mask = rle_[thread_id];
	rle_mask.clear();
	if (has_value_) {
	T value = static_cast<T>(value_[sample_idx]);
	// checking if the value is representable by T, otherwise we
	// just fall back to pick a random pixel.
	if (static_cast<int>(value) == value_[sample_idx]) {
	rle_mask.Init(
	mask, [value](const T &x) { return x == value; });
	}
	} else {
	float threshold = threshold_[sample_idx];
	rle_mask.Init(
	mask, [threshold](const T &x) { return x > threshold; });
	}
	if (rle_mask.count() > 0) {
	auto dist = std::uniform_int_distribution<int64_t>(0, rle_mask.count() - 1);
	flat_idx = rle_mask.find(dist(rng));
	}

[JanuszL]

I would export this as a clear() method.

[JanuszL]

So it is up to the user to make sure that window is inside the image. Maybe it could be mentioned in the documentation - that it is possible and how to deal with that.

[jantonguirao]

Right now this operator is not advertised to be used for cropping. It's a general purpose pixel selector. Do you think we should add this suggestion of usage in the operator documentation?

[JanuszL]

Maybe not if we are not advertising it that way.

[jantonguirao]

I think we could add an example using this operator to calculate a crop center, maybe in one of the existing notebooks

[JanuszL]

Is that print required?

[jantonguirao]

!build

[dali-automaton]

CI MESSAGE: [1796485]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [1796485]: BUILD FAILED

[jantonguirao]

!build

[dali-automaton]

CI MESSAGE: [1796659]: BUILD STARTED

[dali-automaton]

CI MESSAGE: [1796659]: BUILD PASSED