Watermark Detection and Removal on iOS using TensorFlow Lite SSD and OpenCV

This article describes an end‑to‑end pipeline that combines TensorFlow Lite and OpenCV to detect and remove watermarks in images on a mobile app. The workflow covers SSD model training with the TensorFlow Object Detection API, model optimization and conversion to TFLite, inference on iOS using C++, and post‑processing with non‑maximum suppression (NMS) and OpenCV inpainting.

1. SSD Model Training

Clone the TensorFlow models repository and compile the protobuf definitions:

$ cd tensorflow/models
$ protoc object_detection/protos/*.proto --python_out=.

Add the repository to the Python path:

PYTHONPATH=$PYTHONPATH:/your/path/to/tensorflow/models:/your/path/to/tensorflow/models/slim

Prepare labeled data with labelImg , convert XML annotations to CSV, then to TFRecord:

python generate_tfrecord.py --csv_input=test_labels.csv --output_path=test.record

Train the SSD‑MobileNet model (e.g., ssd_mobilenet_v1_coco ) by editing ssd_mobilenet_v1_pets.config (set num_classes and update all PATH_TO_BE_CONFIGURED entries) and running:

python object_detection/train.py \
  --logtostderr \
  --pipeline_config_path=/your/path/training-sets/data-translate/training/ssd_mobilenet_v1_pets.config \
  --train_dir=/your/path/training-sets/data-translate/training

2. Model Optimization and Conversion

After training, strip unnecessary nodes and convert the frozen graph to a TFLite model:

DETECT_PB=$PWD/ssd_mobilenet_v1_coco_2017_11_17/frozen_inference_graph.pb
STRIPPED_PB=$PWD/frozen_inference_graph_stripped.pb
bazel run -c opt tensorflow/contrib/lite/tools/optimize_for_inference \
  -- \
  --input=$DETECT_PB \
  --output=$STRIPPED_PB \
  --frozen_graph=True
bazel run tensorflow/contrib/lite/toco:toco -- \
  --input_file=$STRIPPED_PB \
  --output_file=$DETECT_FB.tflite \
  --input_format=TENSORFLOW_GRAPHDEF \
  --output_format=TFLITE \
  --input_shapes=1,300,300,3 \
  --input_arrays=Preprocessor/sub \
  --output_arrays=concat,concat_1 \
  --inference_type=FLOAT \
  --logtostderr

3. Inference on iOS (C++)

Load an image and preprocess it to the required shape (1×300×300×3) with mean 128 and std 128:

NSString *image_path = FilePathForResourceName(@"test_img", @"jpg");
int image_width, image_height, image_channels;
std::vector<uint8_t> image_data = LoadImageFromFile([image_path UTF8String], ℑ_width, ℑ_height, ℑ_channels);
int wanted_width = 300, wanted_height = 300, wanted_channels = 3;
// ... resize and normalize to float array "in" ...
float *out = interpreter->typed_tensor<float>(input);
// Fill "in" with (pixel - 128)/128 values

Run the interpreter and retrieve the output tensors:

if (interpreter->Invoke() == kTfLiteOk) {
  TfLiteTensor *outputLocations = interpreter->tensor(results[0]);
  TfLiteTensor *outputClasses   = interpreter->tensor(results[1]);
  float *class_data = tflite::GetTensorData<float>(outputClasses);
}

Decode the SSD predictions (center‑size boxes) and apply the sigmoid activation to obtain class scores. Keep detections with score > 0.8 and store classIndex and locationIndex .

4. Non‑Maximum Suppression (NMS)

Sort all boxes by score, select the highest‑scoring box, and discard any remaining box whose IoU with the selected box exceeds a threshold. Repeat until no boxes remain.

5. Watermark Removal with OpenCV

Two approaches are provided:

Inpainting (Telea 2004) : Create a binary mask of the watermark region (same size as the original image) and call cv::inpaint . cv::Mat result; cv::inpaint(original_image, watermark_mask, result, 3, cv::INPAINT_TELEA);

Pixel‑wise color inversion : For semi‑transparent watermarks, compute the original pixel values by reversing the blending operation using the mask.

6. Summary

TensorFlow Lite added SSD support only recently, and documentation is sparse. The iOS C++ example had to be inferred from the Android demo, especially for parsing class outputs and box coordinates. Model conversion steps are sensitive to TensorFlow version changes, so paths and API parameters must be verified for each release.

In practice, watermarks are often placed at fixed positions (corners or center). Future work could add positional priors or attention mechanisms to improve detection efficiency. The current pipeline requires prior knowledge of the watermark style; unknown watermarks cannot be handled. A possible extension is to use GAN‑based image restoration to remove watermarks without explicit detection.

References

https://cloud.google.com/blog/big-data/2017/06/training-an-object-detector-using-cloud-machine-learning-engine

https://github.com/tensorflow/tensorflow/issues/15633#issuecomment-377652630

https://www.pyimagesearch.com/2014/11/17/non-maximum-suppression-object-detection-python/