tool_utils/yolo_world.py

import os 
import logging
import numpy as np
from typing import List
from ultralytics import YOLOWorld

class YoloWorld:
    def __init__(self,model_name = "yolov8x-worldv2.pt"):
        self.model = YOLOWorld(model_name)
        self.model.to(device='cpu')
    
    def run_inference(self,image_path:str,object_prompts:List):
        object_details = []
        self.model.set_classes(object_prompts)
        results = self.model.predict(image_path)
        for result in results:
            for box in result.boxes:
                object_data = {}
                x1, y1, x2, y2 = np.array(box.xyxy.cpu(), dtype=np.int32).squeeze()
                c1,c2  = (x1,y1),(x2,y2)
                confidence = round(float(box.conf.cpu()),2)
                label = f'{results[0].names[int(box.cls)]}' # [{100*round(confidence,2)}%]'
                print("Object Name :{} Bounding Box:{},{} Confidence score {}\n ".format(label ,c1 ,c2,confidence))
                object_data[label] = {
                    'bounding_box':[x1,y1,x2,y2],
                    'confidence':confidence
                    }
                object_details.append(object_data)
        return object_details