Untitled

import json
import base64
import io
from PIL import Image
import yaml
from ultralytics import YOLO
from json import JSONEncoder
import numpy as np


class NumpyFloatValuesEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.float32):
            return float(obj)
        return JSONEncoder.default(self, obj)


def convert_predictions_to_cvat_json(predictions, labels):
    result = []
    if predictions:
        boxes = predictions[0].boxes

        label_ids = boxes.cls.cpu().numpy().tolist()
        confidence_scores = boxes.conf.cpu().numpy().tolist()
        points = predictions[0].masks.xy
        masks = predictions[0].masks.data.cpu().numpy()

        for i in range(len(label_ids)):
            result.append(
                {
                    "confidence": confidence_scores[i],
                    "label": labels.get(label_ids[i], "unknown"),
                    "points": points[i].flatten().tolist(),
                    "mask": masks[i].tolist(),
                    "type": "mask",
                }
            )

    return result


def init_context(context):
    global labels

    context.logger.info("Init context...  0%")

    # Read labels
    with open("/opt/nuclio/function-gpu.yaml", 'rb') as function_file:
        functionconfig = yaml.safe_load(function_file)

    labels_spec = functionconfig['metadata']['annotations']['spec']
    labels = {item['id']: item['name'] for item in json.loads(labels_spec)}

    model = YOLO("/opt/nuclio/model/best.pt")
    context.user_data.model = model

    context.logger.info("Init context...100%")

def handler(context, event):
    context.logger.info("Run Yolo model")
    data = event.body
    buf = io.BytesIO(base64.b64decode(data["image"]))
    threshold = float(data.get("threshold", 0.5))
    image = Image.open(buf)

    predictions = context.user_data.model(image)

    results = convert_predictions_to_cvat_json(predictions, labels)

    # context.logger.info("Results: " + json.dumps(results))
    context.logger.info("Sending results to user...")

    return context.Response(
        body=json.dumps(results),
        headers={},
        content_type='application/json',
        status_code=200
    )