From 394a36886f87fac95957e42c0d55b6d267f07f42 Mon Sep 17 00:00:00 2001
From: srgooglo <srgooglo@ragestudio.net>
Date: Sun, 16 Oct 2022 19:33:20 +0200
Subject: [PATCH] implement download file support

---
 .../src/utils/indecent-prediction/index.js    | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/packages/server/src/utils/indecent-prediction/index.js b/packages/server/src/utils/indecent-prediction/index.js
index cf817671..30fc87bd 100644
--- a/packages/server/src/utils/indecent-prediction/index.js
+++ b/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
 }
\ No newline at end of file