🖼️ 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

ggg4mless
/

RateBooru_Efficient