added possible detection api

README.md

@@ -1,5 +1,5 @@
 ---
-title: Toxity Classifier Api
+title: Image NSFW-Filter API
 emoji: 📉
 colorFrom: pink
 colorTo: yellow

app.py

@@ -1,4 +1,5 @@
 from flask import Flask, request, jsonify, send_from_directory
+from nsfw_detector import predict
 import requests
 import io
 import random
@@ -9,30 +10,55 @@ import string
 import re
 
 app = Flask(__name__)
+model = predict.load_model('./nsfw_mobilenet2.224x224.h5')
 
-API_URL = "https://api-inference.huggingface.co/models/mohsenfayyaz/toxicity-classifier"
-API_TOKEN = os.getenv("HF_READ_TOKEN")  # it is free
-headers = {"Authorization": f"Bearer {API_TOKEN}"}
+@app.route("/imagebuffer", methods=["POST"])
+def check():
+    try:
+        data = request.get_json()
+        imgdata = data["data"]
+
+        #imgdata is like 'data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/2wBD...', save it to a file in ./TEMP
+        imgdata = re.sub('^data:image/.+;base64,', '', imgdata)
+        imgdata = bytes(imgdata, encoding="utf-8")
+        img = Image.open(io.BytesIO(imgdata))
+        filename = ''.join(random.choices(string.ascii_uppercase + string.digits, k=10)) + ".jpg"
+        img.save(f"./TEMP/{filename}")
 
-def query(prompt):
-    payload = {
-        "inputs": prompt
-    }
+        #predict the image
+        result = predict.predict(model, f"./TEMP/{filename}")
+
+        response = {
+            "success": True,
+            "result": result
+        }
+
+    except Exception as e:
+        response = {
+            "success": False,
+            "error": str(e)
+        }
 
-    response = requests.post(API_URL, headers=headers, json=payload)
-    response.raise_for_status()
-    return response.json()
+    return jsonify(response)
 
-@app.route("/generate", methods=["POST"])
-def generate():
+@app.route("/imageurl", methods=["POST"])
+def checkUrl():
     try:
         data = request.get_json()
-        prompt = data["prompt"]
+        url = data["url"]
+
+        #download the image
+        response = requests.get(url)
+        img = Image.open(io.BytesIO(response.content))
+        filename = ''.join(random.choices(string.ascii_uppercase + string.digits, k=10)) + ".jpg"
+        img.save(f"./TEMP/{filename}")
+
+        #predict the image
+        result = predict.predict(model, f"./TEMP/{filename}")
 
-        resp = query(prompt)
         response = {
             "success": True,
-            "result": resp
+            "result": result
         }
 
     except Exception as e:
@@ -43,5 +69,9 @@ def generate():
 
     return jsonify(response)
 
+
 if __name__ == "__main__":
+    if not os.path.exists("./TEMP"):
+        os.makedirs("./TEMP")
+
     app.run(debug=True, host="0.0.0.0", port=7860)

requirements.txt

@@ -1,4 +1,5 @@
 flask
 requests
 pillow
-pycloudflared
+pycloudflared
+nsfw-detector