arthur/xiao
1
0
Fork 0
mirror of https://github.com/arthur-pbty/xiao.git synced 2026-06-03 23:36:43 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix

Browse Source
This commit is contained in:
Dragon Fire
2024-05-08 18:14:08 -04:00
parent 73d69a46d4
commit 059457a61a
1 changed files with 3 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
structures/Tensorflow.js
+3 -5
View File
@@ -73,15 +73,13 @@ module.exports = class Tensorflow {
async stylizeImage(image, styleImg) { async stylizeImage(image, styleImg) {
const imageTensor = await tf.node.decodeImage(image, 3); const imageTensor = await tf.node.decodeImage(image, 3);
const loadedImage = imageTensor.div(tf.scalar(255)).expandDims(); const loadedImage = imageTensor.toFloat().div(tf.scalar(255)).expandDims();
imageTensor.dispose(); imageTensor.dispose();
const styleTensor = tf.node.decodeImage(styleImg, 3); const styleTensor = tf.node.decodeImage(styleImg, 3);
const loadedStyle = styleTensor.div(tf.scalar(255)).expandDims();
styleTensor.dispose(); styleTensor.dispose();
const reshapedStyle = loadedStyle.reshape([1, 1, 1, 100]); const loadedStyle = styleTensor.toFloat().div(tf.scalar(255)).expandDims();
const stylePrediction = await this.styleModel.predict(loadedStyle);
loadedStyle.dispose(); loadedStyle.dispose();
const stylePrediction = await this.styleModel.predict(reshapedStyle);
reshapedStyle.dispose();
const stylizedImage = await this.transformerModel.predict([loadedImage, stylePrediction.squeeze()]); const stylizedImage = await this.transformerModel.predict([loadedImage, stylePrediction.squeeze()]);
loadedImage.dispose(); loadedImage.dispose();
stylePrediction.dispose(); stylePrediction.dispose();