fixed memory leak related to video files

This commit is contained in:
2026-03-29 13:31:27 +02:00
parent 5662f079eb
commit 4cdccb7227
8 changed files with 36 additions and 22 deletions

View File

@@ -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)