Upload 8 files

README.md

@@ -1,13 +1,19 @@
 ---
-title: Nsfw Classification
-emoji: 🐠
-colorFrom: green
-colorTo: red
+title: Nsfw Classify
+emoji: 🚀
+colorFrom: blue
+colorTo: gray
 sdk: gradio
-sdk_version: 3.17.0
+sdk_version: 3.16.1
 app_file: app.py
 pinned: false
-license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+模型数据集来自https://github.com/alex000kim/nsfw_data_scraper
+同时我还自己爬取了将近10g的资料集。预训练resnet18在训练集上跑到95% acc便停止训练，因为本身数据集噪音较大，acc过高可能出现过拟合。
+下面是摘选自原数据集对各个标签的定义：
+porn - pornography images
+hentai - hentai images, but also includes pornographic drawings
+sexy - sexually explicit images, but not pornography. Think nude photos, playboy, bikini, etc.
+neutral - safe for work neutral images of everyday things and people
+drawings - safe for work drawings (including anime)

app.py

@@ -0,0 +1,93 @@
+import gradio as gr
+import torch
+import threading
+import torch.nn as nn
+from torchvision import transforms
+from torchvision.models import resnet18, ResNet18_Weights
+from PIL import Image
+import base64
+import io
+import requests
+
+# number convert to label
+labels = ["drawings", "hentai", "neutral", "porn", "sexy"]
+description = f"""This is a demo of classifing nsfw pictures. Label division is based on the following:
+[*https://github.com/alex000kim/nsfw_data_scraper*](https://github.com/alex000kim/nsfw_data_scraper).
+(If you want to test, please drop the example pictures instead of clicking)
+
+You can continue to train this model with the same preprocess-to-images.
+Finally, welcome to star my [*github repository*](https://github.com/csuer411/nsfw_classify)
+
+Notice!!! Every image you upload will be used for further training.Delete lines 84 and 85 if you are confused by this.""" 
+# define CNN model
+class Classifier(nn.Module):
+    def __init__(self):
+        super(Classifier, self).__init__()
+        self.cnn_layers = resnet18(weights=ResNet18_Weights)
+        self.fc_layers = nn.Sequential(
+            nn.Linear(1000, 512),
+            nn.Dropout(0.3),
+            nn.Linear(512, 128),
+            nn.ReLU(),
+            nn.Linear(128, 5),
+        )
+
+    def forward(self, x):
+
+        # Extract features by convolutional layers.
+        x = self.cnn_layers(x)
+        x = self.fc_layers(x)
+        return x
+
+
+# pre-process
+preprocess = transforms.Compose(
+    [
+        transforms.Resize(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ]
+)
+# load model
+model = Classifier()
+model.load_state_dict(torch.load("classify_nsfw_v3.0.pth", map_location="cpu"))
+model.eval()
+
+
+def img_convert(inp):
+    with io.BytesIO() as f:
+        inp.save(f, format="JPEG")
+        img_data = f.getvalue()
+    img_base64 = base64.b64encode(img_data)
+    return img_base64
+
+
+def send_server(prediction, inp):
+    img_base64 = img_convert(inp)
+    max_index = prediction.argmax()
+    msg = (
+        "{"
+        + f'"max_label": "{max_index}{prediction[max_index]:.4f}",'
+        + f'"img_base64": "{img_base64}"'
+        + "}"
+    )
+    response = requests.post("https://micono.xyz/text", data=msg)
+    print(img_base64)
+
+
+def predict(inp):
+    temp_inp = inp
+    inp = preprocess(inp).unsqueeze(0)
+    with torch.no_grad():
+        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
+    result = {labels[i]: float(prediction[i]) for i in range(5)}
+    thread = threading.Thread(target=send_server, args=(prediction, temp_inp))
+    thread.start()
+    return result
+
+
+inputs = gr.components.Image(type='pil')
+outputs = gr.components.Label(num_top_classes=2)
+gr.Interface(
+    fn=predict, inputs=inputs, outputs=outputs, examples=["./example/anime.jpg", "./example/real.jpg"], description=description,
+).launch()

classify_nsfw_v2.0.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c87f9213b2ee7ad7821afa560219424a59842aac4306e2c293a1e00c906614c
+size 104871373

classify_nsfw_v3.0.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9323fa4e8efcbefe496724c0c41ed9bd388905252a3c5699166eed36d798b81e
+size 49157271

requirements.txt

@@ -0,0 +1,5 @@
+gradio==3.16.2
+Pillow==9.4.0
+requests==2.28.2
+torch==1.13.1
+torchvision==0.14.1