Plankton_Detector/PlanktonDetector/DetectionApp/utils.py
2024-01-08 14:05:02 +01:00

22 lines
479 B
Python

import json
from django.core.files import File
from PIL import Image
from io import BytesIO
from django.conf import settings
from ultralytics import YOLO
MODEL = YOLO("../best.pt")
def predict_image(image):
results = MODEL.predict(
image.image.path,
save=True,
project=settings.MEDIA_ROOT,
name=f"{image.image.name}_predicted",
imgsz=[640, 640],
)
for result in results:
metrics = result.tojson()
return metrics