implement download file support

packages/server/src/utils/indecent-prediction/index.js

@@ -5,6 +5,7 @@ import sharp from "sharp"
 import fs from "fs"
 import path from "path"
 
+import downloadFile from "../download-file"
 import readImage from "../read-image"
 import imageByteArray from "../image-byte-array"
 
@@ -17,10 +18,17 @@ const imageToInput = (image, numChannels) => {
 }
 
 export default async (payload) => {
-    let { image, channels = 3 } = payload
+    let { url, image, channels = 3 } = payload
 
+    let file = null
     const model = await nsfwjs.load()
 
+    if (!image && url) {
+        file = await downloadFile({ url })
+
+        image = file.destination
+    }
+
     // check if image is not a jpg
     if (image.indexOf(".jpg") === -1) {
         // convert image to jpg
@@ -34,6 +42,11 @@ export default async (payload) => {
         fs.writeFileSync(destination, converted)
 
         // set image to the converted image
+        file = {
+            destination,
+            delete: () => fs.unlinkSync(destination),
+        }
+
         image = destination
     }
 
@@ -42,5 +55,9 @@ export default async (payload) => {
 
     const predictions = await model.classify(input)
 
+    if (typeof file.delete === "function") {
+        await file.delete()
+    }
+
     return predictions
 }