test out nsfw

app.py

@@ -4,8 +4,11 @@ import gradio as gr
 from fastcore.all import *
 from fastai.vision.all import *
 import numpy as np
+import opennsfw2 as n2
 import timm
 
+model = n2.make_open_nsfw_model()
+
 
 def parent_labels(o):
     "Label `item` with the parent folder name."
@@ -31,10 +34,23 @@ def predict(image, vtt):
     sprite = PILImage.create(image)
     offsets = []
     images = []
+    tensors = []
+
+    pre_process_data = []
     for left, top, right, bottom in getVTToffsets(vtt):
-        offsets.append((left, top, right, bottom))
         cut_frame = sprite.crop((left, top, left + right, top + bottom))
-        images.append(PILImage.create(np.asarray(cut_frame)))
+        image = n2.preprocess_image(cut_frame, n2.Preprocessing.YAHOO)
+        pre_process_data.append((np.expand_dims(image, axis=0), cut_frame, (left, top, right, bottom)))
+            
+    
+    offsets = []
+    images = []
+    tensors = [i[0] for i in pre_process_data]
+    predictions = model.predict(np.vstack(tensors))
+    for i, prediction in enumerate(predictions):
+        if prediction[0] < 0.5:
+            images.append(PILImage.create(np.asarray(pre_process_data[i][1])))
+            offsets.append(pre_process_data[i][2])
 
     # create dataset
     test_dl = learn.dls.test_dl(images, bs=64)

requirements.txt

@@ -1,4 +1,5 @@
 fastai==2.7.7
-numpy==1.21.5
+numpy==1.24.2
 timm==0.6.7
+opennsfw2==0.10.2
 gradio