ayan4m1/Watermark-Detection-YOLO11-ONNX

Watermark-Detection-YOLO11-ONNX

This is a WebGPU compatible fine-tuning of YOLO11 trained to detect watermarks.

Example

With Transformers.js:

import {
  AutoModel,
  AutoProcessor,
  load_image
} from '@huggingface/transformers';

// require 50% confidence in watermark presence
const threshold = 0.5;
// name of this model
const modelId = 'ayan4m1/Watermark-Detection-YOLO11-ONNX';

// load it using AutoModel and AutoProcessor
const model = await AutoModel.from_pretrained(modelId, { dtype: 'fp32' });
const processor = await AutoProcessor.from_pretrained(modelId);

let watermarked = false;

// load the image and run inference
const image = await load_image(file);
const inputs = await processor(image);
const { output0 } = await model({ images: inputs.pixel_values });

// unpack the results
const permuted = output0[0].transpose(1, 0);
for (const row of permuted.tolist()) {
  // data shape represents a bounding box [xCenter, yCenter, width, height, watermarkProbability]
  const score = row[4];

  if (score < threshold) {
    continue;
  }

  watermarked = true;
  break;
}

if (watermarked) {
  ...
} else {
  ...
}