ImageGenerator Java API
PiperOrigin-RevId: 559310074
This commit is contained in:
MediaPipe Team
Copybara-Service
parent
90781669cb
commit
2ebdb01d43
|
|
@ -0,0 +1,8 @@
|
|||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<manifest xmlns:android="http://schemas.android.com/apk/res/android"
|
||||
package="com.google.mediapipe.tasks.vision.imagegenerator">
|
||||
|
||||
<uses-sdk android:minSdkVersion="24"
|
||||
android:targetSdkVersion="30" />
|
||||
|
||||
</manifest>
|
||||
|
|
@ -0,0 +1,660 @@
|
|||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package com.google.mediapipe.tasks.vision.imagegenerator;
|
||||
|
||||
import android.content.Context;
|
||||
import android.graphics.Bitmap;
|
||||
import android.util.Log;
|
||||
import com.google.auto.value.AutoValue;
|
||||
import com.google.mediapipe.proto.CalculatorOptionsProto.CalculatorOptions;
|
||||
import com.google.mediapipe.framework.AndroidPacketGetter;
|
||||
import com.google.mediapipe.framework.Packet;
|
||||
import com.google.mediapipe.framework.PacketGetter;
|
||||
import com.google.mediapipe.framework.image.BitmapImageBuilder;
|
||||
import com.google.mediapipe.framework.image.MPImage;
|
||||
import com.google.mediapipe.tasks.core.BaseOptions;
|
||||
import com.google.mediapipe.tasks.core.ErrorListener;
|
||||
import com.google.mediapipe.tasks.core.OutputHandler;
|
||||
import com.google.mediapipe.tasks.core.OutputHandler.PureResultListener;
|
||||
import com.google.mediapipe.tasks.core.OutputHandler.ResultListener;
|
||||
import com.google.mediapipe.tasks.core.TaskInfo;
|
||||
import com.google.mediapipe.tasks.core.TaskOptions;
|
||||
import com.google.mediapipe.tasks.core.TaskResult;
|
||||
import com.google.mediapipe.tasks.core.TaskRunner;
|
||||
import com.google.mediapipe.tasks.core.proto.ExternalFileProto;
|
||||
import com.google.mediapipe.tasks.vision.core.BaseVisionTaskApi;
|
||||
import com.google.mediapipe.tasks.vision.core.RunningMode;
|
||||
import com.google.mediapipe.tasks.vision.facelandmarker.FaceLandmarker.FaceLandmarkerOptions;
|
||||
import com.google.mediapipe.tasks.vision.facelandmarker.proto.FaceLandmarkerGraphOptionsProto.FaceLandmarkerGraphOptions;
|
||||
import com.google.mediapipe.tasks.vision.imagegenerator.proto.ConditionedImageGraphOptionsProto.ConditionedImageGraphOptions;
|
||||
import com.google.mediapipe.tasks.vision.imagegenerator.proto.ControlPluginGraphOptionsProto;
|
||||
import com.google.mediapipe.tasks.vision.imagegenerator.proto.ImageGeneratorGraphOptionsProto;
|
||||
import com.google.mediapipe.tasks.vision.imagesegmenter.ImageSegmenter.ImageSegmenterOptions;
|
||||
import com.google.mediapipe.tasks.vision.imagesegmenter.proto.ImageSegmenterGraphOptionsProto.ImageSegmenterGraphOptions;
|
||||
import com.google.protobuf.Any;
|
||||
import com.google.protobuf.ExtensionRegistryLite;
|
||||
import com.google.protobuf.InvalidProtocolBufferException;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import javax.annotation.Nullable;
|
||||
|
||||
/** Performs image generation from a text prompt. */
|
||||
public final class ImageGenerator extends BaseVisionTaskApi {
|
||||
|
||||
private static final String STEPS_STREAM_NAME = "steps";
|
||||
private static final String ITERATION_STREAM_NAME = "iteration";
|
||||
private static final String PROMPT_STREAM_NAME = "prompt";
|
||||
private static final String RAND_SEED_STREAM_NAME = "rand_seed";
|
||||
private static final String SOURCE_CONDITION_IMAGE_STREAM_NAME = "source_condition_image";
|
||||
private static final String CONDITION_IMAGE_STREAM_NAME = "condition_image";
|
||||
private static final String SELECT_STREAM_NAME = "select";
|
||||
private static final int GENERATED_IMAGE_OUT_STREAM_INDEX = 0;
|
||||
private static final int STEPS_OUT_STREAM_INDEX = 1;
|
||||
private static final int ITERATION_OUT_STREAM_INDEX = 2;
|
||||
private static final String TASK_GRAPH_NAME =
|
||||
"mediapipe.tasks.vision.image_generator.ImageGeneratorGraph";
|
||||
private static final String CONDITION_IMAGE_GRAPHS_CONTAINER_NAME =
|
||||
"mediapipe.tasks.vision.image_generator.ConditionedImageGraphContainer";
|
||||
private static final String TAG = "ImageGenerator";
|
||||
private TaskRunner conditionImageGraphsContainerTaskRunner;
|
||||
private Map<ConditionOptions.ConditionType, Integer> conditionTypeIndex;
|
||||
private boolean useConditionImage = false;
|
||||
|
||||
/**
|
||||
* Creates an {@link ImageGenerator} instance from an {@link ImageGeneratorOptions}.
|
||||
*
|
||||
* @param context an Android {@link Context}.
|
||||
* @param generatorOptions an {@link ImageGeneratorOptions} instance.
|
||||
* @throws MediaPipeException if there is an error during {@link ImageGenerator} creation.
|
||||
*/
|
||||
public static ImageGenerator createFromOptions(
|
||||
Context context, ImageGeneratorOptions generatorOptions) {
|
||||
return createFromOptions(context, generatorOptions, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates an {@link ImageGenerator} instance, from {@link ImageGeneratorOptions} and {@link
|
||||
* ConditionOptions}, if plugin models are used to generate an image based on the condition image.
|
||||
*
|
||||
* @param context an Android {@link Context}.
|
||||
* @param generatorOptions an {@link ImageGeneratorOptions} instance.
|
||||
* @param conditionOptions an {@link ConditionOptions} instance.
|
||||
* @throws MediaPipeException if there is an error during {@link ImageGenerator} creation.
|
||||
*/
|
||||
public static ImageGenerator createFromOptions(
|
||||
Context context,
|
||||
ImageGeneratorOptions generatorOptions,
|
||||
@Nullable ConditionOptions conditionOptions) {
|
||||
List<String> inputStreams = new ArrayList<>();
|
||||
inputStreams.addAll(
|
||||
Arrays.asList(
|
||||
"STEPS:" + STEPS_STREAM_NAME,
|
||||
"ITERATION:" + ITERATION_STREAM_NAME,
|
||||
"PROMPT:" + PROMPT_STREAM_NAME,
|
||||
"RAND_SEED:" + RAND_SEED_STREAM_NAME));
|
||||
final boolean useConditionImage = conditionOptions != null;
|
||||
if (useConditionImage) {
|
||||
inputStreams.add("SELECT:" + SELECT_STREAM_NAME);
|
||||
inputStreams.add("CONDITION_IMAGE:" + CONDITION_IMAGE_STREAM_NAME);
|
||||
generatorOptions.conditionOptions = Optional.of(conditionOptions);
|
||||
}
|
||||
List<String> outputStreams =
|
||||
Arrays.asList("IMAGE:image_out", "STEPS:steps_out", "ITERATION:iteration_out");
|
||||
|
||||
OutputHandler<ImageGeneratorResult, Void> handler = new OutputHandler<>();
|
||||
handler.setOutputPacketConverter(
|
||||
new OutputHandler.OutputPacketConverter<ImageGeneratorResult, Void>() {
|
||||
@Override
|
||||
@Nullable
|
||||
public ImageGeneratorResult convertToTaskResult(List<Packet> packets) {
|
||||
int iteration = PacketGetter.getInt32(packets.get(ITERATION_OUT_STREAM_INDEX));
|
||||
int steps = PacketGetter.getInt32(packets.get(STEPS_OUT_STREAM_INDEX));
|
||||
Log.i("ImageGenerator", "Iteration: " + iteration + ", Steps: " + steps);
|
||||
if (iteration != steps - 1) {
|
||||
return null;
|
||||
}
|
||||
Log.i("ImageGenerator", "processing generated image");
|
||||
Packet packet = packets.get(GENERATED_IMAGE_OUT_STREAM_INDEX);
|
||||
Bitmap generatedBitmap = AndroidPacketGetter.getBitmapFromRgb(packet);
|
||||
BitmapImageBuilder bitmapImageBuilder = new BitmapImageBuilder(generatedBitmap);
|
||||
return ImageGeneratorResult.create(
|
||||
bitmapImageBuilder.build(), packet.getTimestamp() / MICROSECONDS_PER_MILLISECOND);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Void convertToTaskInput(List<Packet> packets) {
|
||||
return null;
|
||||
}
|
||||
});
|
||||
handler.setHandleTimestampBoundChanges(true);
|
||||
if (generatorOptions.resultListener().isPresent()) {
|
||||
ResultListener<ImageGeneratorResult, Void> resultListener =
|
||||
new ResultListener<ImageGeneratorResult, Void>() {
|
||||
@Override
|
||||
public void run(ImageGeneratorResult imageGeneratorResult, Void input) {
|
||||
generatorOptions.resultListener().get().run(imageGeneratorResult);
|
||||
}
|
||||
};
|
||||
handler.setResultListener(resultListener);
|
||||
}
|
||||
generatorOptions.errorListener().ifPresent(handler::setErrorListener);
|
||||
TaskRunner runner =
|
||||
TaskRunner.create(
|
||||
context,
|
||||
TaskInfo.<ImageGeneratorOptions>builder()
|
||||
.setTaskName(ImageGenerator.class.getSimpleName())
|
||||
.setTaskRunningModeName(RunningMode.IMAGE.name())
|
||||
.setTaskGraphName(TASK_GRAPH_NAME)
|
||||
.setInputStreams(inputStreams)
|
||||
.setOutputStreams(outputStreams)
|
||||
.setTaskOptions(generatorOptions)
|
||||
.setEnableFlowLimiting(false)
|
||||
.build(),
|
||||
handler);
|
||||
ImageGenerator imageGenerator = new ImageGenerator(runner);
|
||||
if (useConditionImage) {
|
||||
imageGenerator.useConditionImage = true;
|
||||
inputStreams =
|
||||
Arrays.asList(
|
||||
"IMAGE:" + SOURCE_CONDITION_IMAGE_STREAM_NAME, "SELECT:" + SELECT_STREAM_NAME);
|
||||
outputStreams = Arrays.asList("CONDITION_IMAGE:" + CONDITION_IMAGE_STREAM_NAME);
|
||||
OutputHandler<ConditionImageResult, Void> conditionImageHandler = new OutputHandler<>();
|
||||
conditionImageHandler.setOutputPacketConverter(
|
||||
new OutputHandler.OutputPacketConverter<ConditionImageResult, Void>() {
|
||||
@Override
|
||||
public ConditionImageResult convertToTaskResult(List<Packet> packets) {
|
||||
Packet packet = packets.get(0);
|
||||
return new AutoValue_ImageGenerator_ConditionImageResult(
|
||||
new BitmapImageBuilder(AndroidPacketGetter.getBitmapFromRgb(packet)).build(),
|
||||
packet.getTimestamp() / MICROSECONDS_PER_MILLISECOND);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Void convertToTaskInput(List<Packet> packets) {
|
||||
return null;
|
||||
}
|
||||
});
|
||||
conditionImageHandler.setHandleTimestampBoundChanges(true);
|
||||
imageGenerator.conditionImageGraphsContainerTaskRunner =
|
||||
TaskRunner.create(
|
||||
context,
|
||||
TaskInfo.<ImageGeneratorOptions>builder()
|
||||
.setTaskName(ImageGenerator.class.getSimpleName())
|
||||
.setTaskRunningModeName(RunningMode.IMAGE.name())
|
||||
.setTaskGraphName(CONDITION_IMAGE_GRAPHS_CONTAINER_NAME)
|
||||
.setInputStreams(inputStreams)
|
||||
.setOutputStreams(outputStreams)
|
||||
.setTaskOptions(generatorOptions)
|
||||
.setEnableFlowLimiting(false)
|
||||
.build(),
|
||||
conditionImageHandler);
|
||||
imageGenerator.conditionTypeIndex = new HashMap<>();
|
||||
if (conditionOptions.faceConditionOptions().isPresent()) {
|
||||
imageGenerator.conditionTypeIndex.put(
|
||||
ConditionOptions.ConditionType.FACE, imageGenerator.conditionTypeIndex.size());
|
||||
}
|
||||
if (conditionOptions.edgeConditionOptions().isPresent()) {
|
||||
imageGenerator.conditionTypeIndex.put(
|
||||
ConditionOptions.ConditionType.EDGE, imageGenerator.conditionTypeIndex.size());
|
||||
}
|
||||
if (conditionOptions.depthConditionOptions().isPresent()) {
|
||||
imageGenerator.conditionTypeIndex.put(
|
||||
ConditionOptions.ConditionType.DEPTH, imageGenerator.conditionTypeIndex.size());
|
||||
}
|
||||
}
|
||||
return imageGenerator;
|
||||
}
|
||||
|
||||
private ImageGenerator(TaskRunner taskRunner) {
|
||||
super(taskRunner, RunningMode.IMAGE, "", "");
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates an image for iterations and the given random seed. Only valid when the ImageGenerator
|
||||
* is created without condition options.
|
||||
*
|
||||
* @param prompt The text prompt describing the image to be generated.
|
||||
* @param iterations The total iterations to generate the image.
|
||||
* @param seed The random seed used during image generation.
|
||||
*/
|
||||
public ImageGeneratorResult generate(String prompt, int iterations, int seed) {
|
||||
return runIterations(prompt, iterations, seed, null, 0);
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates an image based on the source image for iterations and the given random seed. Only
|
||||
* valid when the ImageGenerator is created with condition options.
|
||||
*
|
||||
* @param prompt The text prompt describing the image to be generated.
|
||||
* @param sourceConditionImage The source image used to create the condition image, which is used
|
||||
* as a guidance for the image generation.
|
||||
* @param conditionType The {@link ConditionOptions.ConditionType} specifying the type of
|
||||
* condition image.
|
||||
* @param iterations The total iterations to generate the image.
|
||||
* @param seed The random seed used during image generation.
|
||||
*/
|
||||
public ImageGeneratorResult generate(
|
||||
String prompt,
|
||||
MPImage sourceConditionImage,
|
||||
ConditionOptions.ConditionType conditionType,
|
||||
int iterations,
|
||||
int seed) {
|
||||
return runIterations(
|
||||
prompt,
|
||||
iterations,
|
||||
seed,
|
||||
createConditionImage(sourceConditionImage, conditionType),
|
||||
conditionTypeIndex.get(conditionType));
|
||||
}
|
||||
|
||||
/**
|
||||
* Create the condition image of specified condition type from the source image. Currently support
|
||||
* face landmarks, depth image and edge image as the condition image.
|
||||
*
|
||||
* @param sourceConditionImage The source image used to create the condition image.
|
||||
* @param conditionType The {@link ConditionOptions.ConditionType} specifying the type of
|
||||
* condition image.
|
||||
*/
|
||||
public MPImage createConditionImage(
|
||||
MPImage sourceConditionImage, ConditionOptions.ConditionType conditionType) {
|
||||
if (!conditionTypeIndex.containsKey(conditionType)) {
|
||||
throw new IllegalArgumentException(
|
||||
"The condition type " + conditionType.name() + " is not created during initialization.");
|
||||
}
|
||||
Map<String, Packet> inputPackets = new HashMap<>();
|
||||
inputPackets.put(
|
||||
SOURCE_CONDITION_IMAGE_STREAM_NAME,
|
||||
conditionImageGraphsContainerTaskRunner
|
||||
.getPacketCreator()
|
||||
.createImage(sourceConditionImage));
|
||||
inputPackets.put(
|
||||
SELECT_STREAM_NAME,
|
||||
conditionImageGraphsContainerTaskRunner
|
||||
.getPacketCreator()
|
||||
.createInt32(conditionTypeIndex.get(conditionType)));
|
||||
ConditionImageResult result =
|
||||
(ConditionImageResult) conditionImageGraphsContainerTaskRunner.process(inputPackets);
|
||||
return result.conditionImage();
|
||||
}
|
||||
|
||||
private ImageGeneratorResult runIterations(
|
||||
String prompt, int steps, int seed, @Nullable MPImage conditionImage, int select) {
|
||||
ImageGeneratorResult result = null;
|
||||
long timestamp = System.currentTimeMillis() * MICROSECONDS_PER_MILLISECOND;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
Map<String, Packet> inputPackets = new HashMap<>();
|
||||
if (i == 0 && useConditionImage) {
|
||||
inputPackets.put(
|
||||
CONDITION_IMAGE_STREAM_NAME, runner.getPacketCreator().createImage(conditionImage));
|
||||
inputPackets.put(SELECT_STREAM_NAME, runner.getPacketCreator().createInt32(select));
|
||||
}
|
||||
inputPackets.put(PROMPT_STREAM_NAME, runner.getPacketCreator().createString(prompt));
|
||||
inputPackets.put(STEPS_STREAM_NAME, runner.getPacketCreator().createInt32(steps));
|
||||
inputPackets.put(ITERATION_STREAM_NAME, runner.getPacketCreator().createInt32(i));
|
||||
inputPackets.put(RAND_SEED_STREAM_NAME, runner.getPacketCreator().createInt32(seed));
|
||||
result = (ImageGeneratorResult) runner.process(inputPackets, timestamp++);
|
||||
}
|
||||
if (useConditionImage) {
|
||||
// Add condition image to the ImageGeneratorResult.
|
||||
return ImageGeneratorResult.create(
|
||||
result.generatedImage(), conditionImage, result.timestampMs());
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
/** Closes and cleans up the task runners. */
|
||||
@Override
|
||||
public void close() {
|
||||
runner.close();
|
||||
conditionImageGraphsContainerTaskRunner.close();
|
||||
}
|
||||
|
||||
/** A container class for the condition image. */
|
||||
@AutoValue
|
||||
protected abstract static class ConditionImageResult implements TaskResult {
|
||||
|
||||
public abstract MPImage conditionImage();
|
||||
|
||||
@Override
|
||||
public abstract long timestampMs();
|
||||
}
|
||||
|
||||
/** Options for setting up an {@link ImageGenerator}. */
|
||||
@AutoValue
|
||||
public abstract static class ImageGeneratorOptions extends TaskOptions {
|
||||
|
||||
/** Builder for {@link ImageGeneratorOptions}. */
|
||||
@AutoValue.Builder
|
||||
public abstract static class Builder {
|
||||
|
||||
/** Sets the text to image model directory storing the model weights. */
|
||||
public abstract Builder setText2ImageModelDirectory(String modelDirectory);
|
||||
|
||||
/** Sets the path to LoRA weights file. */
|
||||
public abstract Builder setLoraWeightsFilePath(String loraWeightsFilePath);
|
||||
|
||||
public abstract Builder setResultListener(
|
||||
PureResultListener<ImageGeneratorResult> resultListener);
|
||||
|
||||
/** Sets an optional {@link ErrorListener}}. */
|
||||
public abstract Builder setErrorListener(ErrorListener value);
|
||||
|
||||
abstract ImageGeneratorOptions autoBuild();
|
||||
|
||||
/** Validates and builds the {@link ImageGeneratorOptions} instance. */
|
||||
public final ImageGeneratorOptions build() {
|
||||
return autoBuild();
|
||||
}
|
||||
}
|
||||
|
||||
abstract String text2ImageModelDirectory();
|
||||
|
||||
abstract Optional<String> loraWeightsFilePath();
|
||||
|
||||
abstract Optional<PureResultListener<ImageGeneratorResult>> resultListener();
|
||||
|
||||
abstract Optional<ErrorListener> errorListener();
|
||||
|
||||
private Optional<ConditionOptions> conditionOptions;
|
||||
|
||||
public static Builder builder() {
|
||||
return new AutoValue_ImageGenerator_ImageGeneratorOptions.Builder()
|
||||
.setText2ImageModelDirectory("");
|
||||
}
|
||||
|
||||
/** Converts an {@link ImageGeneratorOptions} to a {@link Any} protobuf message. */
|
||||
@Override
|
||||
public Any convertToAnyProto() {
|
||||
ImageGeneratorGraphOptionsProto.ImageGeneratorGraphOptions.Builder taskOptionsBuilder =
|
||||
ImageGeneratorGraphOptionsProto.ImageGeneratorGraphOptions.newBuilder();
|
||||
if (conditionOptions != null && conditionOptions.isPresent()) {
|
||||
try {
|
||||
taskOptionsBuilder.mergeFrom(
|
||||
conditionOptions.get().convertToAnyProto().getValue(),
|
||||
ExtensionRegistryLite.getGeneratedRegistry());
|
||||
} catch (InvalidProtocolBufferException e) {
|
||||
Log.e(TAG, "Error converting ConditionOptions to proto. " + e.getMessage());
|
||||
e.printStackTrace();
|
||||
}
|
||||
}
|
||||
taskOptionsBuilder.setText2ImageModelDirectory(text2ImageModelDirectory());
|
||||
if (loraWeightsFilePath().isPresent()) {
|
||||
ExternalFileProto.ExternalFile.Builder externalFileBuilder =
|
||||
ExternalFileProto.ExternalFile.newBuilder();
|
||||
externalFileBuilder.setFileName(loraWeightsFilePath().get());
|
||||
taskOptionsBuilder.setLoraWeightsFile(externalFileBuilder.build());
|
||||
}
|
||||
return Any.newBuilder()
|
||||
.setTypeUrl(
|
||||
"type.googleapis.com/mediapipe.tasks.vision.image_generator.proto.ImageGeneratorGraphOptions")
|
||||
.setValue(taskOptionsBuilder.build().toByteString())
|
||||
.build();
|
||||
}
|
||||
}
|
||||
|
||||
/** Options for setting up the conditions types and the plugin models */
|
||||
@AutoValue
|
||||
public abstract static class ConditionOptions extends TaskOptions {
|
||||
|
||||
/** The supported condition type. */
|
||||
public enum ConditionType {
|
||||
FACE,
|
||||
EDGE,
|
||||
DEPTH
|
||||
}
|
||||
|
||||
/** Builder for {@link ConditionOptions}. At least one type of condition options must be set. */
|
||||
@AutoValue.Builder
|
||||
public abstract static class Builder {
|
||||
public abstract Builder setFaceConditionOptions(FaceConditionOptions faceConditionOptions);
|
||||
|
||||
public abstract Builder setDepthConditionOptions(DepthConditionOptions depthConditionOptions);
|
||||
|
||||
public abstract Builder setEdgeConditionOptions(EdgeConditionOptions edgeConditionOptions);
|
||||
|
||||
abstract ConditionOptions autoBuild();
|
||||
|
||||
/** Validates and builds the {@link ConditionOptions} instance. */
|
||||
public final ConditionOptions build() {
|
||||
ConditionOptions options = autoBuild();
|
||||
if (!options.faceConditionOptions().isPresent()
|
||||
&& !options.depthConditionOptions().isPresent()
|
||||
&& !options.edgeConditionOptions().isPresent()) {
|
||||
throw new IllegalArgumentException(
|
||||
"At least one of `faceConditionOptions`, `depthConditionOptions` and"
|
||||
+ " `edgeConditionOptions` must be set.");
|
||||
}
|
||||
return options;
|
||||
}
|
||||
}
|
||||
|
||||
abstract Optional<FaceConditionOptions> faceConditionOptions();
|
||||
|
||||
abstract Optional<DepthConditionOptions> depthConditionOptions();
|
||||
|
||||
abstract Optional<EdgeConditionOptions> edgeConditionOptions();
|
||||
|
||||
public static Builder builder() {
|
||||
return new AutoValue_ImageGenerator_ConditionOptions.Builder();
|
||||
}
|
||||
|
||||
/**
|
||||
* Converts an {@link ImageGeneratorOptions} to a {@link CalculatorOptions} protobuf message.
|
||||
*/
|
||||
@Override
|
||||
public Any convertToAnyProto() {
|
||||
ImageGeneratorGraphOptionsProto.ImageGeneratorGraphOptions.Builder taskOptionsBuilder =
|
||||
ImageGeneratorGraphOptionsProto.ImageGeneratorGraphOptions.newBuilder();
|
||||
if (faceConditionOptions().isPresent()) {
|
||||
taskOptionsBuilder.addControlPluginGraphsOptions(
|
||||
ControlPluginGraphOptionsProto.ControlPluginGraphOptions.newBuilder()
|
||||
.setBaseOptions(
|
||||
convertBaseOptionsToProto(faceConditionOptions().get().baseOptions()))
|
||||
.setConditionedImageGraphOptions(
|
||||
ConditionedImageGraphOptions.newBuilder()
|
||||
.setFaceConditionTypeOptions(faceConditionOptions().get().convertToProto())
|
||||
.build())
|
||||
.build());
|
||||
}
|
||||
if (edgeConditionOptions().isPresent()) {
|
||||
taskOptionsBuilder.addControlPluginGraphsOptions(
|
||||
ControlPluginGraphOptionsProto.ControlPluginGraphOptions.newBuilder()
|
||||
.setBaseOptions(
|
||||
convertBaseOptionsToProto(edgeConditionOptions().get().baseOptions()))
|
||||
.setConditionedImageGraphOptions(
|
||||
ConditionedImageGraphOptions.newBuilder()
|
||||
.setEdgeConditionTypeOptions(edgeConditionOptions().get().convertToProto())
|
||||
.build())
|
||||
.build());
|
||||
if (depthConditionOptions().isPresent()) {
|
||||
taskOptionsBuilder.addControlPluginGraphsOptions(
|
||||
ControlPluginGraphOptionsProto.ControlPluginGraphOptions.newBuilder()
|
||||
.setBaseOptions(
|
||||
convertBaseOptionsToProto(depthConditionOptions().get().baseOptions()))
|
||||
.setConditionedImageGraphOptions(
|
||||
ConditionedImageGraphOptions.newBuilder()
|
||||
.setDepthConditionTypeOptions(
|
||||
depthConditionOptions().get().convertToProto())
|
||||
.build())
|
||||
.build());
|
||||
}
|
||||
}
|
||||
return Any.newBuilder()
|
||||
.setTypeUrl(
|
||||
"type.googleapis.com/mediapipe.tasks.vision.image_generator.proto.ImageGeneratorGraphOptions")
|
||||
.setValue(taskOptionsBuilder.build().toByteString())
|
||||
.build();
|
||||
}
|
||||
|
||||
/** Options for drawing face landmarks image. */
|
||||
@AutoValue
|
||||
public abstract static class FaceConditionOptions extends TaskOptions {
|
||||
|
||||
/** Builder for {@link FaceConditionOptions}. */
|
||||
@AutoValue.Builder
|
||||
public abstract static class Builder {
|
||||
/** Set the base options for plugin model. */
|
||||
public abstract Builder setBaseOptions(BaseOptions baseOptions);
|
||||
|
||||
/* {@link FaceLandmarkerOptions} used to detect face landmarks in the source image. */
|
||||
public abstract Builder setFaceLandmarkerOptions(
|
||||
FaceLandmarkerOptions faceLandmarkerOptions);
|
||||
|
||||
abstract FaceConditionOptions autoBuild();
|
||||
|
||||
/** Validates and builds the {@link FaceConditionOptions} instance. */
|
||||
public final FaceConditionOptions build() {
|
||||
return autoBuild();
|
||||
}
|
||||
}
|
||||
|
||||
abstract BaseOptions baseOptions();
|
||||
|
||||
abstract FaceLandmarkerOptions faceLandmarkerOptions();
|
||||
|
||||
public static Builder builder() {
|
||||
return new AutoValue_ImageGenerator_ConditionOptions_FaceConditionOptions.Builder();
|
||||
}
|
||||
|
||||
ConditionedImageGraphOptions.FaceConditionTypeOptions convertToProto() {
|
||||
return ConditionedImageGraphOptions.FaceConditionTypeOptions.newBuilder()
|
||||
.setFaceLandmarkerGraphOptions(
|
||||
FaceLandmarkerGraphOptions.newBuilder()
|
||||
.mergeFrom(
|
||||
faceLandmarkerOptions()
|
||||
.convertToCalculatorOptionsProto()
|
||||
.getExtension(FaceLandmarkerGraphOptions.ext))
|
||||
.build())
|
||||
.build();
|
||||
}
|
||||
}
|
||||
|
||||
/** Options for detecting depth image. */
|
||||
@AutoValue
|
||||
public abstract static class DepthConditionOptions extends TaskOptions {
|
||||
|
||||
/** Builder for {@link DepthConditionOptions}. */
|
||||
@AutoValue.Builder
|
||||
public abstract static class Builder {
|
||||
|
||||
/** Set the base options for plugin model. */
|
||||
public abstract Builder setBaseOptions(BaseOptions baseOptions);
|
||||
|
||||
/** {@link ImageSegmenterOptions} used to detect depth image from the source image. */
|
||||
public abstract Builder setImageSegmenterOptions(
|
||||
ImageSegmenterOptions imageSegmenterOptions);
|
||||
|
||||
abstract DepthConditionOptions autoBuild();
|
||||
|
||||
/** Validates and builds the {@link DepthConditionOptions} instance. */
|
||||
public final DepthConditionOptions build() {
|
||||
DepthConditionOptions options = autoBuild();
|
||||
return options;
|
||||
}
|
||||
}
|
||||
|
||||
abstract BaseOptions baseOptions();
|
||||
|
||||
abstract ImageSegmenterOptions imageSegmenterOptions();
|
||||
|
||||
public static Builder builder() {
|
||||
return new AutoValue_ImageGenerator_ConditionOptions_DepthConditionOptions.Builder();
|
||||
}
|
||||
|
||||
ConditionedImageGraphOptions.DepthConditionTypeOptions convertToProto() {
|
||||
return ConditionedImageGraphOptions.DepthConditionTypeOptions.newBuilder()
|
||||
.setImageSegmenterGraphOptions(
|
||||
imageSegmenterOptions()
|
||||
.convertToCalculatorOptionsProto()
|
||||
.getExtension(ImageSegmenterGraphOptions.ext))
|
||||
.build();
|
||||
}
|
||||
}
|
||||
|
||||
/** Options for detecting edge image. */
|
||||
@AutoValue
|
||||
public abstract static class EdgeConditionOptions {
|
||||
|
||||
/**
|
||||
* Builder for {@link EdgeConditionOptions}.
|
||||
*
|
||||
* <p>These parameters are used to config Canny edge algorithm of OpenCV.
|
||||
*
|
||||
* <p>See more details:
|
||||
* https://docs.opencv.org/3.4/dd/d1a/group__imgproc__feature.html#ga04723e007ed888ddf11d9ba04e2232de
|
||||
*/
|
||||
@AutoValue.Builder
|
||||
public abstract static class Builder {
|
||||
|
||||
/** Set the base options for plugin model. */
|
||||
public abstract Builder setBaseOptions(BaseOptions baseOptions);
|
||||
|
||||
/** First threshold for the hysteresis procedure. */
|
||||
public abstract Builder setThreshold1(Float threshold1);
|
||||
|
||||
/** Second threshold for the hysteresis procedure. */
|
||||
public abstract Builder setThreshold2(Float threshold2);
|
||||
|
||||
/** Aperture size for the Sobel operator. Typical range is 3~7. */
|
||||
public abstract Builder setApertureSize(Integer apertureSize);
|
||||
|
||||
/**
|
||||
* flag, indicating whether a more accurate L2 norm should be used to calculate the image
|
||||
* gradient magnitude ( L2gradient=true ), or whether the default L1 norm is enough (
|
||||
* L2gradient=false ).
|
||||
*/
|
||||
public abstract Builder setL2Gradient(Boolean l2Gradient);
|
||||
|
||||
abstract EdgeConditionOptions autoBuild();
|
||||
|
||||
/** Validates and builds the {@link EdgeConditionOptions} instance. */
|
||||
public final EdgeConditionOptions build() {
|
||||
return autoBuild();
|
||||
}
|
||||
}
|
||||
|
||||
abstract BaseOptions baseOptions();
|
||||
|
||||
abstract Float threshold1();
|
||||
|
||||
abstract Float threshold2();
|
||||
|
||||
abstract Integer apertureSize();
|
||||
|
||||
abstract Boolean l2Gradient();
|
||||
|
||||
public static Builder builder() {
|
||||
return new AutoValue_ImageGenerator_ConditionOptions_EdgeConditionOptions.Builder()
|
||||
.setThreshold1(100f)
|
||||
.setThreshold2(200f)
|
||||
.setApertureSize(3)
|
||||
.setL2Gradient(false);
|
||||
}
|
||||
|
||||
ConditionedImageGraphOptions.EdgeConditionTypeOptions convertToProto() {
|
||||
return ConditionedImageGraphOptions.EdgeConditionTypeOptions.newBuilder()
|
||||
.setThreshold1(threshold1())
|
||||
.setThreshold2(threshold2())
|
||||
.setApertureSize(apertureSize())
|
||||
.setL2Gradient(l2Gradient())
|
||||
.build();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,44 @@
|
|||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package com.google.mediapipe.tasks.vision.imagegenerator;
|
||||
|
||||
import com.google.auto.value.AutoValue;
|
||||
import com.google.mediapipe.framework.image.MPImage;
|
||||
import com.google.mediapipe.tasks.core.TaskResult;
|
||||
import java.util.Optional;
|
||||
|
||||
/** Represents the image generation results generated by {@link ImageGenerator}. */
|
||||
@AutoValue
|
||||
public abstract class ImageGeneratorResult implements TaskResult {
|
||||
|
||||
/** Create an {@link ImageGeneratorResult} instance from the generated image. */
|
||||
public static ImageGeneratorResult create(
|
||||
MPImage generatedImage, MPImage conditionImage, long timestampMs) {
|
||||
return new AutoValue_ImageGeneratorResult(
|
||||
generatedImage, Optional.of(conditionImage), timestampMs);
|
||||
}
|
||||
|
||||
/** Create an {@link ImageGeneratorResult} instance from the generated image. */
|
||||
public static ImageGeneratorResult create(MPImage generatedImage, long timestampMs) {
|
||||
return new AutoValue_ImageGeneratorResult(generatedImage, Optional.empty(), timestampMs);
|
||||
}
|
||||
|
||||
public abstract MPImage generatedImage();
|
||||
|
||||
public abstract Optional<MPImage> conditionImage();
|
||||
|
||||
@Override
|
||||
public abstract long timestampMs();
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user