|
|
--- |
|
|
library_name: transformers.js |
|
|
license: agpl-3.0 |
|
|
pipeline_tag: object-detection |
|
|
datasets: |
|
|
- bastienp/visible-watermark-pita |
|
|
base_model: |
|
|
- Ultralytics/YOLO11 |
|
|
base_model_relation: finetune |
|
|
tags: |
|
|
- watermark |
|
|
--- |
|
|
|
|
|
# Watermark-Detection-YOLO11-ONNX |
|
|
|
|
|
This is a WebGPU compatible fine-tuning of YOLO11 trained to detect watermarks. |
|
|
|
|
|
## Example |
|
|
|
|
|
With Transformers.js: |
|
|
|
|
|
```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 { |
|
|
... |
|
|
} |
|
|
``` |