Class: NSFW::Model

Inherits:
Object
  • Object
show all
Defined in:
lib/nsfw/model.rb

Constant Summary collapse

MODEL_PATH = 
"vendor/onnx_models/nsfw.onnx"
CATEGORIES = 
['drawings', 'hentai', 'neutral', 'porn', 'sexy']
SAFETY_THRESHOLD = 
0.85

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(lazy: false) ⇒ Model

Returns a new instance of Model.



11
12
13
 
# File 'lib/nsfw/model.rb', line 11

def initialize(lazy: false)
  load_model! unless lazy
end

Instance Attribute Details

#model ⇒ Object (readonly)

Returns the value of attribute model.



9
10
11
 
# File 'lib/nsfw/model.rb', line 9

def model
  @model
end

Instance Method Details

#loaded? ⇒ Boolean

Returns:

  • (Boolean) 


25
26
27
 
# File 'lib/nsfw/model.rb', line 25

def loaded?
  !@model.nil?
end

#predict(tensor) ⇒ Object 



15
16
17
18
 
# File 'lib/nsfw/model.rb', line 15

def predict(tensor)
  prediction = make_prediction(tensor)
  format_prediction(prediction)
end

#reshape_tensor(tensor) ⇒ Object 



29
30
31
 
# File 'lib/nsfw/model.rb', line 29

def reshape_tensor(tensor)
  OnnxRuntime::Utils.reshape(tensor, [1, 224, 224, 3])
end

#safe?(image) ⇒ Boolean

Returns:

  • (Boolean) 


20
21
22
23
 
# File 'lib/nsfw/model.rb', line 20

def safe?(image)
  prediction = predict(image.tensor)
  prediction["neutral"] >= SAFETY_THRESHOLD
end