Spaces:
Running
Running
| import os | |
| import logging | |
| import cv2 | |
| import numpy as np | |
| from typing import List | |
| import torch | |
| import random | |
| 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 | |
| def draw_bboxes(rgb_frame,boxes,labels,color=None,line_thickness=3): | |
| rgb_frame = cv2.imread(rgb_frame) | |
| rgb_frame = cv2.cvtColor(rgb_frame,cv2.COLOR_BGR2RGB) | |
| tl = line_thickness or round(0.002 * (rgb_frame.shape[0] + rgb_frame.shape[1]) / 2) + 1 # line/font thickness | |
| rgb_frame_copy = rgb_frame.copy() | |
| if color is None : | |
| color = color or [random.randint(0, 255) for _ in range(3)] | |
| for box,label in zip(boxes,labels): | |
| if box.type() == 'torch.IntTensor': | |
| box = box.numpy() | |
| # extract coordinates | |
| x1,y1,x2,y2 = box | |
| c1,c2 = (x1,y1),(x2,y2) | |
| # Draw rectangle | |
| cv2.rectangle(rgb_frame_copy, c1,c2, color, thickness=tl, lineType=cv2.LINE_AA) | |
| tf = max(tl - 1, 1) # font thickness | |
| # label = label2id[int(label.numpy())] | |
| t_size = cv2.getTextSize(str(label), 0, fontScale=tl / 3, thickness=tf)[0] | |
| c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 | |
| cv2.putText(rgb_frame_copy, str(label), (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) | |
| return rgb_frame_copy | |
| def run_yolo_infer(self,image_path:str,object_prompts:List): | |
| self.model.set_classes(object_prompts) | |
| results = self.model.predict(image_path) | |
| processed_predictions = [] | |
| bounding_boxes = [] | |
| labels = [] | |
| scores = [] | |
| for result in results: | |
| for i,box in enumerate(result.boxes): | |
| x1, y1, x2, y2 = np.array(box.xyxy.cpu(), dtype=np.int32).squeeze() | |
| bounding_boxes.append([x1,y1,x2,y2]) | |
| labels.append(int(box.cls.cpu())) | |
| scores.append(round(float(box.conf.cpu()),2)) | |
| processed_predictions.append(dict(boxes= torch.tensor(bounding_boxes), | |
| labels= torch.IntTensor(labels), | |
| scores=torch.tensor(scores)) | |
| ) | |
| detected_image = self.draw_bboxes(rgb_frame=image_path, | |
| boxes=processed_predictions[0]['boxes'], | |
| labels=processed_predictions[0]['labels'] | |
| ) | |
| cv2.imwrite('final_mask.png',detected_image) | |
| return "Predicted image mask : final_mask.jpg . Details :{}".format(processed_predictions[0]) | |