You've already forked JapariArchive
fixed memory leak related to video files
This commit is contained in:
@@ -4,6 +4,8 @@ from io import BytesIO
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from transformers import ViTForImageClassification, ViTImageProcessor
|
||||
|
||||
from Database.x_classes import HavoxLabel
|
||||
|
||||
class VoxClassifier():
|
||||
def __init__(self):
|
||||
self.feature_extractor = ViTImageProcessor.from_pretrained('Classifier/HAV0X/')
|
||||
@@ -21,9 +23,13 @@ class VoxClassifier():
|
||||
return self.model.config.id2label[predicted_class_idx]
|
||||
except Exception as ex:
|
||||
print(ex)
|
||||
return "Rejected"
|
||||
return HavoxLabel.Rejected
|
||||
|
||||
async def classify_async(self, image_bytes):
|
||||
if not image_bytes:
|
||||
#mainly for video files
|
||||
print("image bytes was null")
|
||||
return HavoxLabel.Rejected
|
||||
with ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(self.classify, image_bytes)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user