tensorflow_nsfw_load.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('nsfw_classifier_tokenizer.pickle', 'rb') as f:\n",
    "    tokenizer = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('nsfw_classifier.pickle', 'rb') as f:\n",
    "    model = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the vocabulary size and embedding dimensions\n",
    "vocab_size = 10000\n",
    "embedding_dim = 64\n",
    "\n",
    "# Pad the prompt and negative prompt sequences\n",
    "max_sequence_length = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "def preprocess(text, isfirst = True):\n",
    "    if isfirst:\n",
    "        if type(text) == str: pass\n",
    "        elif type(text) == list:\n",
    "            output = []\n",
    "            for i in text:\n",
    "                output.append(preprocess(i))\n",
    "            return(output)\n",
    "            \n",
    "\n",
    "    text = re.sub('<.*?>', '', text)\n",
    "    text = re.sub('\\(+', '(', text)\n",
    "    text = re.sub('\\)+', ')', text)\n",
    "    matchs = re.findall('\\(.*?\\)', text)\n",
    "    \n",
    "    for _ in matchs:\n",
    "        text = text.replace(_, preprocess(_[1:-1], isfirst=False) )\n",
    "\n",
    "    text = text.replace('\\n', ',').replace('|',',')\n",
    "\n",
    "    if isfirst: \n",
    "        output = text.split(',')\n",
    "        output = list(map(lambda x: x.strip(), output))\n",
    "        output = [x for x in output if x != '']\n",
    "        return ', '.join(output)\n",
    "        # return output\n",
    "\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def postprocess(prompts, negative_prompts, outputs, print_percentage = True):\n",
    "    for idx, i in enumerate(prompts):\n",
    "        print('*****************************************************************')\n",
    "        if print_percentage:\n",
    "            print(f\"prompt: {i}\\nnegative_prompt: {negative_prompts[idx]}\\npredict: {outputs[idx][0]} --{outputs[idx][1]}%\")\n",
    "        else:\n",
    "            print(f\"prompt: {i}\\nnegative_prompt: {negative_prompts[idx]}\\npredict: {outputs[idx][0]}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make predictions on new data\n",
    "prompt = [\"a landscape with trees and mountains in the background\", 'nude, sexy, 1girl, nsfw']\n",
    "negative_prompt = [\"nsfw\", 'worst quality']\n",
    "\n",
    "x_new = tokenizer.texts_to_sequences( preprocess(prompt) )\n",
    "z_new = tokenizer.texts_to_sequences( preprocess(negative_prompt) )\n",
    "x_new = tf.keras.preprocessing.sequence.pad_sequences(x_new, maxlen=max_sequence_length)\n",
    "z_new = tf.keras.preprocessing.sequence.pad_sequences(z_new, maxlen=max_sequence_length)\n",
    "y_new = model.predict([x_new, z_new])\n",
    "y_new = list(map(lambda x:(\"NSFW\", float(\"{:.2f}\".format(x[0]*100)) ) if x[0]>0.5 else (\"SFW\", float(\"{:.2f}\".format(100-x[0]*100))), y_new))\n",
    "\n",
    "\n",
    "print(\"Prediction:\", y_new)\n",
    "postprocess(prompt, negative_prompt, y_new, print_percentage=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}