✨ add anime pictures and real pictures to be intelligently classified

src/store/modules/uploadStore.ts

@@ -4,7 +4,8 @@ import i18n from "@/locales";
 import {NSFWJS} from "nsfwjs";
 import {message} from "ant-design-vue";
 
-import {loadCocoSsd, loadMobileNet} from "@/utils/tfjs";
+// import {loadCocoSsd, loadMobileNet} from "@/utils/tfjs/tfjs.ts";
+import {loadModel, predictImage} from "@/utils/tfjs/anime_classifier.ts";
 
 export const useUploadStore = defineStore(
     'upload',
@@ -25,11 +26,19 @@ export const useUploadStore = defineStore(
                 message.error(i18n.global.t('comment.illegalImage'));
                 return false;
             }
-            const predictions = await loadMobileNet(image);
-            console.log(predictions);
+            // const predictions = await loadMobileNet(image);
+            // console.log(predictions);
+            //
+            // const prediction = await loadCocoSsd(image);
+            // console.log(prediction);
 
-            const prediction = await loadCocoSsd(image);
-            console.log(prediction);
+            const model = await loadModel('/tfjs/anime_classifier/model.json');
+
+            // 进行预测
+            const output = await predictImage(model, image);
+            console.log(output);
+
+            // console.log('Predicted Class:', predictedClass);
             return true;
         }
 

src/utils/tfjs/anime_classifier.ts

@@ -0,0 +1,23 @@
+import * as tf from '@tensorflow/tfjs';
+
+// 封装处理图像和推理的工具函数
+export async function loadModel(modelPath) {
+    const model = await tf.loadGraphModel(modelPath);
+    console.log('Model Loaded');
+    return model;
+}
+
+//   处理图片并进行推理
+export async function predictImage(model, imageElement) {
+    // 将图片转换为张量
+    const tensor = tf.browser.fromPixels(imageElement).toFloat();
+    const resized = tf.image.resizeBilinear(tensor, [224, 224]);  // 调整图片大小为模型输入大小
+    const input = resized.expandDims(0);  // 增加批次维度
+
+    // 进行推理
+    const prediction = model.predict(input);
+
+    // 获取预测结果并返回
+    const resultArray = await prediction.array();
+    return resultArray[0]; // 返回第一项的预测结果
+}