talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update java image segmenter to always output confidence masks and optionally output category mask.

Browse Source 
PiperOrigin-RevId: 521804641
This commit is contained in:
MediaPipe Team
2023-04-04 10:41:59 -07:00
committed by Copybara-Service
parent 7c2930102d
commit 33cad24a5a
5 changed files with 107 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/imagesegmenter/ImageSegmenterTest.java
+33 -46
View File
@@ -61,14 +61,13 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CATEGORY_MASK)
.setOutputCategoryMask(true)
.build();
ImageSegmenter imageSegmenter =
ImageSegmenter.createFromOptions(ApplicationProvider.getApplicationContext(), options);
ImageSegmenterResult actualResult = imageSegmenter.segment(getImageFromAsset(inputImageName));
List<MPImage> segmentations = actualResult.segmentations();
assertThat(segmentations.size()).isEqualTo(1);
MPImage actualMaskBuffer = actualResult.segmentations().get(0);
assertThat(actualResult.categoryMask().isPresent()).isTrue();
MPImage actualMaskBuffer = actualResult.categoryMask().get();
MPImage expectedMaskBuffer = getImageFromAsset(goldenImageName);
verifyCategoryMask(
actualMaskBuffer, expectedMaskBuffer, GOLDEN_MASK_SIMILARITY, MAGNIFICATION_FACTOR);
@@ -81,15 +80,14 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.build();
ImageSegmenter imageSegmenter =
ImageSegmenter.createFromOptions(ApplicationProvider.getApplicationContext(), options);
ImageSegmenterResult actualResult = imageSegmenter.segment(getImageFromAsset(inputImageName));
List<MPImage> segmentations = actualResult.segmentations();
List<MPImage> segmentations = actualResult.confidenceMasks();
assertThat(segmentations.size()).isEqualTo(21);
// Cat category index 8.
MPImage actualMaskBuffer = actualResult.segmentations().get(8);
MPImage actualMaskBuffer = segmentations.get(8);
MPImage expectedMaskBuffer = getImageFromAsset(goldenImageName);
verifyConfidenceMask(actualMaskBuffer, expectedMaskBuffer, GOLDEN_MASK_SIMILARITY);
}
@@ -102,40 +100,36 @@ public class ImageSegmenterTest {
ImageSegmenterOptions.builder()
.setBaseOptions(
BaseOptions.builder().setModelAssetPath(SELFIE_128x128_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.build();
ImageSegmenter imageSegmenter =
ImageSegmenter.createFromOptions(ApplicationProvider.getApplicationContext(), options);
ImageSegmenterResult actualResult = imageSegmenter.segment(getImageFromAsset(inputImageName));
List<MPImage> segmentations = actualResult.segmentations();
List<MPImage> segmentations = actualResult.confidenceMasks();
assertThat(segmentations.size()).isEqualTo(2);
// Selfie category index 1.
MPImage actualMaskBuffer = actualResult.segmentations().get(1);
MPImage actualMaskBuffer = segmentations.get(1);
MPImage expectedMaskBuffer = getImageFromAsset(goldenImageName);
verifyConfidenceMask(actualMaskBuffer, expectedMaskBuffer, GOLDEN_MASK_SIMILARITY);
}
// TODO: enable this unit test once activation option is supported in metadata.
// @Test
// public void segment_successWith144x256Segmentation() throws Exception {
// final String inputImageName = "mozart_square.jpg";
// final String goldenImageName = "selfie_segm_144_256_3_expected_mask.jpg";
// ImageSegmenterOptions options =
// ImageSegmenterOptions.builder()
// .setBaseOptions(
// BaseOptions.builder().setModelAssetPath(SELFIE_144x256_MODEL_FILE).build())
// .setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
// .build();
// ImageSegmenter imageSegmenter =
// ImageSegmenter.createFromOptions(ApplicationProvider.getApplicationContext(), options);
// ImageSegmenterResult actualResult =
// imageSegmenter.segment(getImageFromAsset(inputImageName));
// List<MPImage> segmentations = actualResult.segmentations();
// assertThat(segmentations.size()).isEqualTo(1);
// MPImage actualMaskBuffer = actualResult.segmentations().get(0);
// MPImage expectedMaskBuffer = getImageFromAsset(goldenImageName);
// verifyConfidenceMask(actualMaskBuffer, expectedMaskBuffer, GOLDEN_MASK_SIMILARITY);
// }
@Test
public void segment_successWith144x256Segmentation() throws Exception {
final String inputImageName = "mozart_square.jpg";
final String goldenImageName = "selfie_segm_144_256_3_expected_mask.jpg";
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(
BaseOptions.builder().setModelAssetPath(SELFIE_144x256_MODEL_FILE).build())
.build();
ImageSegmenter imageSegmenter =
ImageSegmenter.createFromOptions(ApplicationProvider.getApplicationContext(), options);
ImageSegmenterResult actualResult = imageSegmenter.segment(getImageFromAsset(inputImageName));
List<MPImage> segmentations = actualResult.confidenceMasks();
assertThat(segmentations.size()).isEqualTo(1);
MPImage actualMaskBuffer = segmentations.get(0);
MPImage expectedMaskBuffer = getImageFromAsset(goldenImageName);
verifyConfidenceMask(actualMaskBuffer, expectedMaskBuffer, GOLDEN_MASK_SIMILARITY);
}
@Test
public void getLabels_success() throws Exception {
@@ -165,7 +159,6 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.build();
ImageSegmenter imageSegmenter =
ImageSegmenter.createFromOptions(ApplicationProvider.getApplicationContext(), options);
@@ -287,16 +280,15 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.setRunningMode(RunningMode.IMAGE)
.build();
ImageSegmenter imageSegmenter =
ImageSegmenter.createFromOptions(ApplicationProvider.getApplicationContext(), options);
ImageSegmenterResult actualResult = imageSegmenter.segment(getImageFromAsset(inputImageName));
List<MPImage> segmentations = actualResult.segmentations();
List<MPImage> segmentations = actualResult.confidenceMasks();
assertThat(segmentations.size()).isEqualTo(21);
// Cat category index 8.
MPImage actualMaskBuffer = actualResult.segmentations().get(8);
MPImage actualMaskBuffer = segmentations.get(8);
MPImage expectedMaskBuffer = getImageFromAsset(goldenImageName);
verifyConfidenceMask(actualMaskBuffer, expectedMaskBuffer, GOLDEN_MASK_SIMILARITY);
}
@@ -309,12 +301,11 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.setRunningMode(RunningMode.IMAGE)
.setResultListener(
(segmenterResult, inputImage) -> {
verifyConfidenceMask(
segmenterResult.segmentations().get(8),
segmenterResult.confidenceMasks().get(8),
expectedResult,
GOLDEN_MASK_SIMILARITY);
})
@@ -331,7 +322,6 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.setRunningMode(RunningMode.VIDEO)
.build();
ImageSegmenter imageSegmenter =
@@ -341,10 +331,10 @@ public class ImageSegmenterTest {
ImageSegmenterResult actualResult =
imageSegmenter.segmentForVideo(
getImageFromAsset(inputImageName), /* timestampsMs= */ i);
List<MPImage> segmentations = actualResult.segmentations();
List<MPImage> segmentations = actualResult.confidenceMasks();
assertThat(segmentations.size()).isEqualTo(21);
// Cat category index 8.
MPImage actualMaskBuffer = actualResult.segmentations().get(8);
MPImage actualMaskBuffer = segmentations.get(8);
verifyConfidenceMask(actualMaskBuffer, expectedMaskBuffer, GOLDEN_MASK_SIMILARITY);
}
}
@@ -357,12 +347,11 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.setRunningMode(RunningMode.VIDEO)
.setResultListener(
(segmenterResult, inputImage) -> {
verifyConfidenceMask(
segmenterResult.segmentations().get(8),
segmenterResult.confidenceMasks().get(8),
expectedResult,
GOLDEN_MASK_SIMILARITY);
})
@@ -384,12 +373,11 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.setRunningMode(RunningMode.LIVE_STREAM)
.setResultListener(
(segmenterResult, inputImage) -> {
verifyConfidenceMask(
segmenterResult.segmentations().get(8),
segmenterResult.confidenceMasks().get(8),
expectedResult,
GOLDEN_MASK_SIMILARITY);
})
@@ -411,12 +399,11 @@ public class ImageSegmenterTest {
ImageSegmenterOptions options =
ImageSegmenterOptions.builder()
.setBaseOptions(BaseOptions.builder().setModelAssetPath(DEEPLAB_MODEL_FILE).build())
.setOutputType(ImageSegmenterOptions.OutputType.CONFIDENCE_MASK)
.setRunningMode(RunningMode.LIVE_STREAM)
.setResultListener(
(segmenterResult, inputImage) -> {
verifyConfidenceMask(
segmenterResult.segmentations().get(8),
segmenterResult.confidenceMasks().get(8),
expectedResult,
GOLDEN_MASK_SIMILARITY);
})
mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/interactivesegmenter/InteractiveSegmenterTest.java
+5 -2
View File
@@ -60,7 +60,10 @@ public class InteractiveSegmenterTest {
ApplicationProvider.getApplicationContext(), options);
MPImage image = getImageFromAsset(inputImageName);
ImageSegmenterResult actualResult = imageSegmenter.segment(image, roi);
List<MPImage> segmentations = actualResult.segmentations();
// TODO update to correct category mask output.
// After InteractiveSegmenter updated according to (b/276519300), update this to use
// categoryMask field instead of confidenceMasks.
List<MPImage> segmentations = actualResult.confidenceMasks();
assertThat(segmentations.size()).isEqualTo(1);
}
@@ -79,7 +82,7 @@ public class InteractiveSegmenterTest {
ApplicationProvider.getApplicationContext(), options);
ImageSegmenterResult actualResult =
imageSegmenter.segment(getImageFromAsset(inputImageName), roi);
List<MPImage> segmentations = actualResult.segmentations();
List<MPImage> segmentations = actualResult.confidenceMasks();
assertThat(segmentations.size()).isEqualTo(2);
}
}