Spaces:
Running
Running
File size: 5,080 Bytes
d283c04 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | import { describe, it, expect } from 'vitest';
import { embedWatermark } from '../core/embedder.js';
import { detectWatermark, detectWatermarkMultiFrame } from '../core/detector.js';
import { getPreset } from '../core/presets.js';
import type { PresetName } from '../core/types.js';
/**
* Crop a Y plane by removing pixels from the edges.
*/
function cropYPlane(
yPlane: Uint8Array,
width: number,
height: number,
left: number,
top: number,
right: number,
bottom: number,
): { cropped: Uint8Array; croppedW: number; croppedH: number } {
const croppedW = width - left - right;
const croppedH = height - top - bottom;
const cropped = new Uint8Array(croppedW * croppedH);
for (let y = 0; y < croppedH; y++) {
for (let x = 0; x < croppedW; x++) {
cropped[y * croppedW + x] = yPlane[(y + top) * width + (x + left)];
}
}
return { cropped, croppedW, croppedH };
}
/**
* Generate a synthetic Y plane (gradient + noise)
*/
function generateTestFrame(width: number, height: number): Uint8Array {
const frame = new Uint8Array(width * height);
for (let y = 0; y < height; y++) {
for (let x = 0; x < width; x++) {
const gradient = ((x + y) / (width + height)) * 200 + 20;
const noise = (Math.random() - 0.5) * 20;
frame[y * width + x] = Math.max(0, Math.min(255, Math.round(gradient + noise)));
}
}
return frame;
}
describe('Embed → Detect Round-trip', () => {
const width = 512;
const height = 512;
const payload = new Uint8Array([0xDE, 0xAD, 0xBE, 0xEF]);
const key = 'test-secret-key-2024';
// Test the presets that should work with a 512x512 frame
for (const presetName of ['light', 'moderate', 'strong'] as PresetName[]) {
it(`should round-trip with ${presetName} preset`, () => {
const config = getPreset(presetName);
const frame = generateTestFrame(width, height);
// Embed
const embedResult = embedWatermark(frame, width, height, payload, key, config);
expect(embedResult.psnr).toBeGreaterThan(15); // Higher presets sacrifice invisibility for robustness
// Detect
const detectResult = detectWatermark(embedResult.yPlane, width, height, key, config);
expect(detectResult.detected).toBe(true);
expect(detectResult.payload).not.toBeNull();
expect(Array.from(detectResult.payload!)).toEqual(Array.from(payload));
});
}
it('should not detect watermark with wrong key', () => {
const config = getPreset('moderate');
const frame = generateTestFrame(width, height);
const embedResult = embedWatermark(frame, width, height, payload, key, config);
const detectResult = detectWatermark(embedResult.yPlane, width, height, 'wrong-key', config);
expect(detectResult.detected).toBe(false);
});
it('should not detect watermark on unwatermarked frame', () => {
const config = getPreset('moderate');
const frame = generateTestFrame(width, height);
const detectResult = detectWatermark(frame, width, height, key, config);
expect(detectResult.detected).toBe(false);
});
});
describe('Crop-resilient detection', () => {
const width = 512;
const height = 512;
const payload = new Uint8Array([0xDE, 0xAD, 0xBE, 0xEF]);
const key = 'test-secret-key-2024';
it('should detect after arbitrary crop with multiple frames', () => {
const config = getPreset('strong');
// Generate multiple distinct frames (simulating video — at least 32 frames)
const frames: Uint8Array[] = [];
for (let i = 0; i < 32; i++) {
const frame = generateTestFrame(width, height);
const embedResult = embedWatermark(frame, width, height, payload, key, config);
// Crop by arbitrary offset (breaks DWT pixel pairing)
const { cropped } = cropYPlane(
embedResult.yPlane, width, height, 7, 3, 5, 9
);
frames.push(cropped);
}
const croppedW = width - 7 - 5;
const croppedH = height - 3 - 9;
const detectResult = detectWatermarkMultiFrame(
frames, croppedW, croppedH, key, config, { cropResilient: true }
);
expect(detectResult.detected).toBe(true);
expect(detectResult.payload).not.toBeNull();
expect(Array.from(detectResult.payload!)).toEqual(Array.from(payload));
});
it('should detect after large crop with multiple frames', () => {
const config = getPreset('strong');
const frames: Uint8Array[] = [];
for (let i = 0; i < 32; i++) {
const frame = generateTestFrame(width, height);
const embedResult = embedWatermark(frame, width, height, payload, key, config);
const { cropped } = cropYPlane(
embedResult.yPlane, width, height, 50, 50, 50, 50
);
frames.push(cropped);
}
const croppedW = width - 100;
const croppedH = height - 100;
const detectResult = detectWatermarkMultiFrame(
frames, croppedW, croppedH, key, config, { cropResilient: true }
);
expect(detectResult.detected).toBe(true);
expect(detectResult.payload).not.toBeNull();
expect(Array.from(detectResult.payload!)).toEqual(Array.from(payload));
});
});
|