COCO_1K/ls_backend/yolo_ls.py

import dagshub
import mlflow
import os

from label_studio_ml.model import LabelStudioMLBase
from ultralytics import YOLO

class YoloLS(LabelStudioMLBase):

    def __init__(self, **kwargs):
        super(YoloLS, self).__init__(**kwargs)
        
        # pre-initialize your variables here
        from_name, schema = list(self.parsed_label_config.items())[0]
        self.from_name = from_name
        self.to_name = schema['to_name'][0]
        self.labels = schema['labels']

        self.user = os.getenv("DAGSHUB_USER_NAME")
        self.token = os.getenv("DAGSHUB_TOKEN")
        self.repo = os.getenv("DAGSHUB_REPO_NAME")
        self.host = os.getenv("DAGSHUB_CLIENT_HOST")
        
        dagshub.auth.add_app_token(token=self.token, host=self.host)
            
        dagshub.init(repo_name=self.repo, repo_owner=self.user)
        
        self.model = YOLO('yolov8n-seg.pt', task='segment')
#         client = mlflow.MlflowClient()
#         name = 'YOLOv8 Custom'
#         version = client.get_latest_versions(name=name)[0].version
        
#         self.model_version = f'{name}:{version}'

#         model_uri = f'models:/{name}/{version}'
        
#         self.model = mlflow.pyfunc.load_model(model_uri)
    
    def image_uri_to_https(self, uri):
        if uri.startswith('http'):
            return uri
        elif uri.startswith('repo://'):
            link_data = uri.split("repo://")[-1].split("/")
            commit, tree_path = link_data[0], "/".join(link_data[1:])
            return f"{self.host}/api/v1/repos/{self.user}/{self.repo}/raw/{commit}/{tree_path}"
        elif uri.startswith('/'):
            return f"{self.host}{uri}"
        raise FileNotFoundError(f'Unknown URI {uri}')
        
    def predict(self, tasks, **kwargs):
        """ This is where inference happens: 
            model returns the list of predictions based on input list of tasks
            
            :param tasks: Label Studio tasks in JSON format
        """
        results = []
        
        for task in tasks:
            uri = task['data']['image']
            url = self.image_uri_to_https(uri)
            
            if self.host != "https://dagshub.com":
                url = url.replace("https://",f"https://{self.user}:{self.token}@")
            
            preds = self.model.predict(url)[0]
        
            lowest_conf = 2.0
        
            img_results = []
            
            boxes = preds.boxes.cpu().numpy()
            masks = preds.masks
            
            for i in range(len(boxes.cls)):
                conf = float(boxes.conf[i])
                if conf < lowest_conf:
                    lowest_conf = conf
                
                img_results.append({
                    'type': 'polygonlabels',
                    'to_name': self.to_name,
                    'from_name': self.from_name,
                    'image_rotation': 0,
                    'original_height': preds.orig_shape[0],
                    'original_width': preds.orig_shape[1],
                    'value': {
                        'closed': True,
                        'points': (masks.xyn[i] * 100).astype(float).tolist(),
                        'polygonlabels': [self.labels[int(boxes.cls[i])]],
                    },
                    'score': conf
                })
                
                # img_results = {
                #     'type': 'polygonlabels',
                #     'to_name': self.to_name,
                #     'from_name': self.from_name,
                #     'image_rotation': 0,
                #     'original_height': preds.orig_shape[0],
                #     'original_width': preds.orig_shape[1],
                #     'value': {
                #         'closed': True,
                #         'points': (masks.xyn[i] * 100).astype(float).tolist(),
                #         'polygonlabels': self.labels[int(boxes.cls[i])],
                #     },
                #     'score': conf
                # }

            results.append({
                'result': img_results,
                'score': lowest_conf
            })
            
        return results

    def fit(self, event, data, **kwargs):
        """ This is where training happens: train your model given list of completions,
            then returns dict with created links and resources

            :param completions: aka annotations, the labeling results from Label Studio 
            :param workdir: current working directory for ML backend
        """
        # save some training outputs to the job result
        return {'random': 1}