目標檢測系列（六）labelstudio實現自動化標注

一、啟用圖片文件服務

用Nginx啟用圖片服務，配置好映射路徑。

新建圖片文件夾，將文件夾下的圖片路徑存儲到txt文件中

訪問地址（文件夾）：http://112.12.19.122:8081/urls/ml-backend-test/

進入labelstudio將txt文件路徑填入，點擊Add URL將圖片導入項目進行標注。

二、啟用模型服務

首先pip安裝label-studio-ml

進入到projects文件夾，將init_model.py放入該文件夾，然后執行命令label-studio-ml init my_backend來初始化模型文件夾。init_model.py的代碼如下：

#!/user/bin/env?python3
#?-*-?coding:?utf-8?-*-
from label_studio_ml.model import LabelStudioMLBaseclass DummyModel(LabelStudioMLBase):def __init__(self, **kwargs):# don't forget to call base class constructorsuper(DummyModel, self).__init__(**kwargs)# you can preinitialize variables with keys needed to extract info from tasks and annotations and form predictionsfrom_name, schema = list(self.parsed_label_config.items())[0]self.from_name = from_nameself.to_name = schema['to_name'][0]self.labels = schema['labels']def predict(self, tasks, **kwargs):""" This is where inference happens: model returnsthe list of predictions based on input list of tasks"""predictions = []for task in tasks:predictions.append({'score': 0.987,? # prediction overall score, visible in the data manager columns'model_version': 'delorean-20151021',? # all predictions will be differentiated by model version'result': [{'from_name': self.from_name,'to_name': self.to_name,'type': 'choices','score': 0.5,? # per-region score, visible in the editor'value': {'choices': [self.labels[0]]}}]})return predictionsdef fit(self, annotations, **kwargs):""" This is where training happens: train your model given list of annotations,then returns dict with created links and resources"""return {'path/to/created/model': 'my/model.bin'}

進入到my_backend文件夾，可以看到下述文件：

在my_backend文件夾下新建model文件夾，將訓練好的YOLO模型文件放入model下：

修改my_backend文件夾下的model.py，代碼如下：

#!/user/bin/env?python3# ?-*-?coding:?utf-8?-*-import osfrom typing import List, Dict, Optionalimport torchfrom label_studio_ml.model import LabelStudioMLBasefrom label_studio_ml.utils import get_single_tag_keys, get_local_pathimport loggingfrom ultralytics import YOLOfrom PIL import Image# 設置日志logger = logging.getLogger(__name__)logging.basicConfig(level=logging.INFO)MODEL_PATH = os.getenv('MODEL_PATH', '/data/projects/my_ml_backend/model/best.pt')DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')class DummyModel(LabelStudioMLBase):"""Custom ML Backend model"""def __init__(self, **kwargs):super(DummyModel, self).__init__(**kwargs)from_name, schema = list(self.parsed_label_config.items())[0]self.from_name = from_nameself.to_name = schema['to_name'][0]self.labels = schema['labels']# 訓練參數self.train_epochs = int(os.getenv('TRAIN_EPOCHS', 150))self.train_batch_size = int(os.getenv('TRAIN_BATCH_SIZE', 18))def predict(self, tasks: List[Dict], context: Optional[Dict] = None, **kwargs):task = tasks[0]print(f'''\Run prediction on {tasks}Received context: {context}Project ID: {task['id']}Label config: {self.label_config}Parsed JSON Label config: {self.parsed_label_config}''')img_url = task['data']['image']image_path = self.get_local_path(url=img_url)print(f'image_path: {image_path}')# Getting URL and loading imageimage = Image.open(image_path)# Height and width of imageoriginal_width, original_height = image.size# Creating list for predictions and variable for scorespredictions = []scores = 0i = 0# Initialize self variablesself.from_name, self.to_name, self.value, self.classes = get_single_tag_keys(self.parsed_label_config, 'RectangleLabels', 'Image')# 加載自己的yolov11模型logger.info(f"加載YOLO11模型: {MODEL_PATH}")self.model = YOLO(MODEL_PATH)# 檢查GPU可用性logger.info(f"使用設備: {'GPU ?' if DEVICE.type == 'cuda' else 'CPU ??'}")# 改動的地方, 增加了conf配置, 只有conf>=0.5的才會被標記出來# 默認conf是0.25, 不改的話被標注的地方肯能會很多, 根據自己的實際情況配置# Getting prediction using modelresults = self.model.predict(image, conf=0.5)# print(results)# Getting mask segments, boxes from model predictionfor result in results:for i, prediction in enumerate(result.boxes):score = prediction.conf.item()label_index = int(prediction.cls.item())xyxy = prediction.xyxy[0].tolist()# print(f"{i} prediction", prediction)# x_center, y_center, w, h = boxpredictions.append({"id": str(i),"from_name": self.from_name,"to_name": self.to_name,"type": "rectanglelabels","score": score,"original_width": original_width,"original_height": original_height,"image_rotation": 0,"value": {"rotation": 0,# 坐標轉換, 只有轉換后才能標注在正確的位置"x": xyxy[0] / original_width * 100,"y": xyxy[1] / original_height * 100,"width": (xyxy[2] - xyxy[0]) / original_width * 100,"height": (xyxy[3] - xyxy[1]) / original_height * 100,"rectanglelabels": [self.labels[label_index]]}})scores += scorelogger.info(f"預測完成: 檢測到 {len(predictions)} 個對象")# Dict with final dicts with predictionsfinal_prediction = [{"result": predictions,"score": scores / (i + 1),"model_version": "11x"}]return final_predictiondef fit(self, event, data, **kwargs):"""使用新標注數據訓練模型參數:event: 事件類型 ('ANNOTATION_CREATED', 'ANNOTATION_UPDATED')data: 包含標注數據的字典**kwargs: 額外參數"""# 檢查是否有訓練數據if not self.train_output:logger.info("初始化訓練數據存儲")self.train_output = {'image_paths': [],'labels': []}# 獲取標注信息annotation = data['annotation']image_url = annotation['task']['data']['image']image_path = self.get_local_path(image_url)# 解析標注結果bboxes = []for result in annotation['result']:if result['from_name'] == self.from_name:value = result['value']label = value['rectanglelabels'][0]# 獲取圖像尺寸image = Image.open(image_path)img_width, img_height = image.size# 轉換為絕對坐標x = value['x'] * img_width / 100y = value['y'] * img_height / 100width = value['width'] * img_width / 100height = value['height'] * img_height / 100# YOLO格式: [class_idx, x_center, y_center, width, height] (歸一化)x_center = (x + width / 2) / img_widthy_center = (y + height / 2) / img_heightnorm_width = width / img_widthnorm_height = height / img_heightclass_idx = self.labels.index(label)bboxes.append([class_idx, x_center, y_center, norm_width, norm_height])# 保存訓練數據self.train_output['image_paths'].append(image_path)self.train_output['labels'].append(bboxes)logger.info(f"收到新標注: 圖像={image_path}, 標注數={len(bboxes)}")logger.info(f"當前訓練集大小: {len(self.train_output['image_paths'])}")# 當有足夠數據時開始訓練if len(self.train_output['image_paths']) >= 10:logger.info("達到最小訓練集大小，開始訓練...")self.train_model()# 重置訓練數據self.train_output = {'image_paths': [],'labels': []}# 返回新模型信息return {'model_path': MODEL_PATH,'model_version': f"retrained-{len(self.train_output['image_paths'])}"}return {}def train_model(self):"""使用收集的標注數據訓練模型"""logger.info("準備訓練數據...")# 創建YOLO格式的訓練數據目錄結構train_dir = 'yolo_train_data'images_dir = os.path.join(train_dir, 'images')labels_dir = os.path.join(train_dir, 'labels')os.makedirs(images_dir, exist_ok=True)os.makedirs(labels_dir, exist_ok=True)# 創建數據集描述文件with open(os.path.join(train_dir, 'dataset.yaml'), 'w') as f:f.write(f"train: {os.path.abspath(images_dir)}\n")f.write(f"nc: {len(self.labels)}\n")f.write(f"names: {self.labels}\n")# 準備訓練數據for i, (image_path, bboxes) in enumerate(zip(self.train_output['image_paths'],self.train_output['labels'])):# 復制圖像img = Image.open(image_path)img_filename = f'train_{i}.jpg'img.save(os.path.join(images_dir, img_filename))# 創建標簽文件label_filename = f'train_{i}.txt'with open(os.path.join(labels_dir, label_filename), 'w') as f:for bbox in bboxes:class_idx, x_center, y_center, width, height = bboxf.write(f"{class_idx} {x_center} {y_center} {width} {height}\n")logger.info(f"訓練數據準備完成: {len(self.train_output['image_paths'])} 張圖像")# 訓練模型 (這里簡化了實際訓練過程)logger.info(f"開始訓練模型 (模擬) - 周期={self.train_epochs}, 批次大小={self.train_batch_size}")# 調用YOLO的訓練腳本:# import subprocess# subprocess.run(['python', 'train.py'])# 參數配置： --img 640 --batch {self.train_batch_size} --epochs {self.train_epochs}#????????? --data {os.path.join(train_dir, 'dataset.yaml')} --weights {MODEL_PATH}logger.info("訓練完成! 模型已更新")# 重新加載訓練后的模型# self.model = YOLO(MODEL_PATH)

最后執行命令label-studio-ml start my_ml_backend -p 9094來啟動模型后端服務