forked from prehistoric-systems/comixify
47e5ea3231
* Add popularity
* Update popularity model to dockerfile
* Fix dockerfile
* Revert "Fix dockerfile"
This reverts commit ea24fba604.
* Fix dockerfile
* Fix filename
* Fix filename
* Fix mkdir
* Fix mkdir
* Fix #9
* Zip popularity model
* Fix model filename
* Remove comixify volume
* Fix minimum number of segments
* Add debug print
* Update model 1 file
16 lines
504 B
Python
16 lines
504 B
Python
import pickle
|
|
import os.path
|
|
|
|
MODEL_PATH = 'popularity/pretrained_model/svr_test_11.10.sk'
|
|
|
|
|
|
class PopularityPredictor:
|
|
def __init__(self):
|
|
if not os.path.exists(MODEL_PATH):
|
|
print("Model file does not exist.")
|
|
with open(MODEL_PATH, 'rb') as fp:
|
|
self.svr = pickle.load(fp, encoding='latin1')
|
|
|
|
def get_popularity_score(self, image_feature):
|
|
image_feature = image_feature.reshape(1, -1)
|
|
return self.svr.predict(image_feature)
|