a

zeruniverse · Nov 17, 2024 · a31a159 · a31a159
1 parent b52c320
commit a31a159
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 10 deletions.
diff --git a/dependencies/models/robustdecoder.onnx b/dependencies/models/robustdecoder.onnx
diff --git a/dependencies/models/encoder.onnx → dependencies/models/robustencoder.onnx b/dependencies/models/encoder.onnx → dependencies/models/robustencoder.onnx
diff --git a/dependencies/models/decoder.onnx → dependencies/models/visualdecoder.onnx b/dependencies/models/decoder.onnx → dependencies/models/visualdecoder.onnx
diff --git a/dependencies/models/visualencoder.onnx b/dependencies/models/visualencoder.onnx
diff --git a/stego.js b/stego.js
@@ -2,8 +2,10 @@
 
 import * as ort from './dependencies/onnxruntime-web-1.20.0/ort.webgpu.min.mjs';
 
-let encoderSession = null;
-let decoderSession = null;
+let encoderSessionv = null;
+let decoderSessionv = null;
+let encoderSessionr = null;
+let decoderSessionr = null;
 let wasmModule = null;
 
 // Initialize Codecs
@@ -14,16 +16,18 @@ export async function initCodecs() {
     // Load models
     try {
         wasmModule = await createWasmModule();
-        encoderSession = await ort.InferenceSession.create('./dependencies/models/encoder.onnx', sessionOptions);
-        decoderSession = await ort.InferenceSession.create('./dependencies/models/decoder.onnx', sessionOptions);
+        encoderSessionv = await ort.InferenceSession.create('./dependencies/models/visualencoder.onnx', sessionOptions);
+        decoderSessionv = await ort.InferenceSession.create('./dependencies/models/visualdecoder.onnx', sessionOptions);
+        encoderSessionr = await ort.InferenceSession.create('./dependencies/models/robustencoder.onnx', sessionOptions);
+        decoderSessionr = await ort.InferenceSession.create('./dependencies/models/robustdecoder.onnx', sessionOptions);
     } catch (e) {
         throw new Error('Error loading models or WASM module: ' + e.message);
     }
 }
 
 // Check if codecs are ready
 export function isCodecsReady() {
-    return encoderSession !== null && decoderSession !== null && wasmModule !== null;
+    return encoderSessionr !== null && decoderSessionr !== null && wasmModule !== null;
 }
 
 // Load the Emscripten WASM module
@@ -84,7 +88,7 @@ export async function loadIMGtoCanvas(inputid, canvasid, maxsize = 0) {
 }
 
 // Write message to canvas
-export async function writeMsgToCanvas(canvasid, msg, password = '', check_valid = true) {
+export async function writeMsgToCanvas(canvasid, msg, password = '', model_type = 0, check_valid = true) {
     try {
         const canvas = document.getElementById(canvasid);
         if (!canvas) {
@@ -119,7 +123,7 @@ export async function writeMsgToCanvas(canvasid, msg, password = '', check_valid
         };
 
         // Run encoder model
-        const results = await encoderSession.run(feeds);
+        const results = await ((model_type==0)?encoderSessionv.run(feeds):encoderSessionr.run(feeds));
         const encodedImg = results['encoded_img'].data;
 
         // Draw encoded image back to canvas
@@ -130,7 +134,7 @@ export async function writeMsgToCanvas(canvasid, msg, password = '', check_valid
         if (check_valid) {
             // Read the message back to verify
             try{
-                const decodedMsg = await readMsgFromCanvas(canvasid, password);
+                const decodedMsg = await readMsgFromCanvas(canvasid, password, model_type);
                 if (decodedMsg !== msg) {
                     return false;
                 }
@@ -145,7 +149,7 @@ export async function writeMsgToCanvas(canvasid, msg, password = '', check_valid
 }
 
 // Read message from canvas
-export async function readMsgFromCanvas(canvasid, password = '') {
+export async function readMsgFromCanvas(canvasid, password = '', model_type=0) {
     try {
         const canvas = document.getElementById(canvasid);
         if (!canvas) {
@@ -166,7 +170,7 @@ export async function readMsgFromCanvas(canvasid, password = '') {
         };
 
         // Run decoder model
-        const results = await decoderSession.run(feeds);
+        const results = await ((model_type==0)?decoderSessionv.run(feeds):decoderSessionr.run(feeds));
         const dataOutput = results['data'].data; // Float32Array of length 65536
 
         // Use the WASM module to decode dataOutput to bytes