ggg4mless/RateBooru_Efficient

🖼️ RateBooru Efficient - Danbooru Rating Classifier

This is a simple image classification model fine-tuned from EfficientNetB1 to classify Danbooru-style image ratings into 3 categories:

• general
• questionable
• explicit

🎯 Purpose

This model was created as a learning project and also intended for basic NSFW moderation tasks based on Danbooru rating style.

📊 Dataset

• Total images: 6,000
  • Training: 4,800
  • Validation: 1,200
• All images are tagged with rating metadata from Danbooru.
• The dataset is private / custom-collected, not provided here.

🧠 Model Details

• Base model: EfficientNetB1 from TensorFlow/Keras
• Fine-tuned for 50 epochs
• Input size: 240x240 RGB
• Output: Softmax classification (3 classes)

📈 Final Training Metrics

Metric	Training	Validation
Accuracy	75.24%	70.42%
Loss	0.5916	0.6452
Precision	79.07%	72.93%
Recall	69.09%	67.58%
Learning Rate	2e-6	-

🚀 Usage (TensorFlow)

import tensorflow as tf
import numpy as np
from PIL import Image
import os

# === Configuration ===
MODEL_PATH = 'ratebooru_efficientnetb1.keras'
IMAGE_PATH = 'example.jpg'  # Change this to your image file path
IMG_SIZE = (240, 240)
CLASS_NAMES = ['explicit', 'general', 'questionable']

# === Load & Preprocess Image ===
def preprocess_image(image_path):
    try:
        img = Image.open(image_path).convert('RGB')
        img = img.resize(IMG_SIZE)
        img_array = tf.keras.utils.img_to_array(img)
        return tf.expand_dims(img_array, 0)
    except Exception as e:
        print(f"Error loading image: {e}")
        return None

# === Predict ===
def predict(model, image_tensor):
    predictions = model.predict(image_tensor)
    score = tf.nn.softmax(predictions[0])
    return {CLASS_NAMES[i]: float(score[i]) for i in range(len(CLASS_NAMES))}

# === Main ===
if not os.path.exists(MODEL_PATH):
    print(f"Model not found at '{MODEL_PATH}'")
    exit()

if not os.path.exists(IMAGE_PATH):
    print(f"Image not found at '{IMAGE_PATH}'")
    exit()

print("Loading model...")
model = tf.keras.models.load_model(MODEL_PATH)
print("Model loaded.")

print(f"Predicting image: {IMAGE_PATH}")
image_tensor = preprocess_image(IMAGE_PATH)
if image_tensor is None:
    exit()

results = predict(model, image_tensor)

print("\\n--- Prediction Result ---")
for class_name, confidence in sorted(results.items(), key=lambda x: x[1], reverse=True):
    print(f"{class_name:<15}: {confidence * 100:.2f}%")

Loading model...
Model loaded.
Predicting image: example.png

--- Prediction Result ---
general        : 54.99%
questionable   : 22.57%
explicit       : 22.44%

⚠️ Disclaimer

• This model is not perfect and might misclassify borderline content.
• Do not use it for serious moderation or legal filtering without thorough evaluation.
• Dataset is based on subjective human tagging (Danbooru), which may include biases.

📄 License

MIT License