Tổng quan về nghiên cứu
Trong lĩnh vực thị giác máy tính hiện đại, kỹ thuật phân đoạn hình ảnh đóng vai trò quan trọng trong nhiều ứng dụng thực tế. Đặc biệt trong nông nghiệp, lâm nghiệp và giám sát hệ sinh thái, khả năng nhận diện và phân đoạn chính xác các bộ phận của thực vật mang lại giá trị to lớn trong việc nâng cao hiệu suất sản xuất và theo dõi biến đổi môi trường. Cành cây là thành phần thiết yếu của thực vật, trực tiếp ảnh hưởng đến trạng thái sinh trưởng, sức khỏe và đặc điểm phân bố của cây, từ đó tác động đến sự ổn định của toàn hệ sinh thái.
Dòng mô hình YOLO (You Only Look Once) đã trở thành lựa chọn phổ biến trong các tác vụ phát hiện đối tượng nhờ hiệu suất vượt trội. YOLOv8, phiên bản mới nhất của dòng model này, kết hợp những tiến bộ trong học sâu và thị giác máy tính, mang lại độ chính xác cao hơn và tốc độ xử lý nhanh hơn. Tuy nhiên, khi áp dụng cho các tác vụ đặc thù như phân đoạn hình ảnh cành cây, YOLOv8 gốc vẫn còn những hạn chế nhất định.
Nghiên cứu này xây dựng hệ thống phân đoạn hình ảnh cành cây hiệu quả dựa trên YOLOv8 được cải tiến, sử dụng bộ dữ liệu "new treebranch 3" bao gồm 2000 hình ảnh cành cây đa dạng về hình thái và bối cảnh. Thiết kế đơn lớp duy nhất cho phép mô hình tập trung hoàn toàn vào việc trích xuất đặc điểm của cành cây, tránh nhiễu từ các lớp đối tượng khác.
Thông tin bộ dữ liệu
Bộ dữ liệu "new treebranch 3" được xây dựng với mục tiêu huấn luyện và cải tiến mô hình YOLOv8-seg cho tác vụ phân đoạn hình ảnh cành cây. Tập dữ liệu tập trung vào một lớp đối tượng duy nhất có tên "Branch", giúp mô hình tập trung vào đặc điểm riêng biệt của cành cây mà không bị phân tán bởi các lớp khác.
Mặc dù chỉ có một lớp duy nhất, các hình ảnh trong bộ dữ liệu đã được tuyển chọn kỹ lưỡng để đảm bao phủ đa dạng về môi trường, điều kiện ánh sáng và hình dạng cành cây. Quá trình gán nhãn tuân thủ nghiêm ngặt các tiêu chuẩn chất lượng cao, đảm bảo mỗi vùng cành cây trong hình ảnh đều được phân đoạn chính xác.
Để tăng cường khả năng tổng quát hóa của mô hình, các kỹ thuật tăng cường dữ liệu đã được áp dụng, bao gồm xoay, phóng to/thu nhỏ, cắt裁 và điều chỉnh màu sắc. Chiến lược này không chỉ mở rộng quy mô dữ liệu huấn luyện mà còn giúp mô hình thích ứng tốt hơn với các hình ảnh chưa từng thấy.
Cấu trúc mã nguồn cốt lõi
Dưới đây là các thành phần chính của mô hình dự đoán, được giải thích chi tiết:
class SegmentPredictor:
"""
Lớp cơ sở cho bộ dự đoán phân đoạn hình ảnh.
Thuộc tính:
config (Namespace): Cấu hình cho bộ dự đoán.
output_path (Path): Thư mục lưu kết quả.
neural_model (nn.Module): Mô hình mạng nơ-ron sử dụng để dự đoán.
compute_device (torch.device): Thiết bị tính toán.
data_loader (Dataset): Tập dữ liệu cho dự đoán.
"""
def __init__(self, config_path=DEFAULT_CFG, param_override=None, callback_list=None):
"""
Khởi tạo SegmentPredictor.
Tham số:
config_path (str): Đường dẫn file cấu hình.
param_override (dict): Ghi đè tham số cấu hình.
callback_list: Danh sách các callback tùy chỉnh.
"""
self.config = load_config(config_path, param_override)
self.output_path = create_output_directory(self.config)
self.warmup_completed = False
self.neural_model = None
self.compute_device = None
self.data_loader = None
self.callback_list = callback_list or default_callbacks()
def prepare_input(self, img_data):
"""
Chuẩn bị hình ảnh đầu vào trước khi suy luận.
Tham số:
img_data (torch.Tensor | List[np.ndarray]): Hình ảnh đầu vào.
Trả về:
torch.Tensor: Hình ảnh đã được xử lý.
"""
is_not_tensor = not isinstance(img_data, torch.Tensor)
if is_not_tensor:
img_data = np.concatenate(self.apply_transform(img_data))
img_data = img_data[..., ::-1].transpose((0, 3, 1, 2))
img_data = np.ascontiguousarray(img_data)
img_data = torch.from_numpy(img_data)
img_data = img_data.to(self.compute_device)
img_data = img_data.half() if self.neural_model.fp16 else img_data.float()
if is_not_tensor:
img_data /= 255.0
return img_data
def run_inference(self, input_data, *args, **kwargs):
"""Thực hiện suy luận trên hình ảnh đã xử lý."""
return self.neural_model(input_data, augment=self.config.augment)
def process_stream(self, source=None, model=None, *args, **kwargs):
"""Xử lý luồng video từ camera và lưu kết quả."""
if not self.neural_model:
self.initialize_model(model)
self.configure_source(source if source is not None else self.config.source)
for batch_data in self.data_loader:
img_path, original_img, video_cap, status = batch_data
processed_img = self.prepare_input(original_img)
predictions = self.run_inference(processed_img, *args, **kwargs)
self.prediction_results = self.format_output(predictions, processed_img, original_img)
for idx in range(len(original_img)):
self.save_results(idx, self.prediction_results, (img_path[idx], processed_img, original_img[idx]))
yield from self.prediction_results
def initialize_model(self, model_weight, verbose=True):
"""Khởi tạo mô hình YOLO và chuyển sang chế độ đánh giá."""
self.neural_model = AutoBackend(model_weight or self.config.model,
device=select_device(self.config.device))
self.compute_device = self.neural_model.device
self.neural_model.eval()
def save_results(self, index, results, batch_info):
"""Lưu kết quả suy luận ra file hoặc thư mục."""
file_path, img_tensor, _ = batch_info
result_item = results[index]
if self.config.save or self.config.show:
self.visualization = result_item.plot()
if self.config.save_txt:
result_item.save_txt(f'{self.output_path / "labels" / file_path.stem}.txt')Giải thích các thành phần
Class SegmentPredictor đóng vai trò là lớp nền tảng cho bộ dự đoán YOLO, cung cấp các phương thức khởi tạo, tiền xử lý, suy luận và xử lý luồng dữ liệu. Phương thức __init__ thiết lập các thuộc tính ban đầu bao gồm cấu hình, đường dẫn lưu trữ, mô hình và thiết bị tính toán.
Phương thức prepare_input chịu trách nhiệm chuyển đổi hình ảnh đầu vào sang định dạng phù hợp với yêu cầu của mô hình, bao gồm chuyển đổi không gian màu và chuẩn hóa giá trị pixel. Phương thức run_inference thực hiện quá trình suy luận trên hình ảnh đã được xử lý.
Phương thức process_stream cho phép xử lý thời gian thực từ các nguồn video như camera hoặc file video, thực hiện suy luận trên từng khung hình và lưu kết quả. Phương thức initialize_model khởi tạo mô hình YOLO và đặt nó ở chế độ đánh giá, trong khi save_results ghi kết quả ra file hoặc hiển thị trực quan.
Kiến trúc hệ thống
Dự án Ultralytics YOLO được tổ chức thành nhiều module chức năng riêng biệt, đảm bảo quy trình hoàn chỉnh từ xử lý dữ liệu, huấn luyện mô hình đến đánh giá và dự đoán kết quả.
Module xử lý dữ liệu thực hiện tiền xử lý và quản lý file đầu vào, đảm bảo dữ liệu đáp ứng yêu cầu của mô hình. Module dự đoán cung cấp chức năng suy luận cốt lõi, hỗ trợ đa dạng nguồn đầu vào như hình ảnh, video và camera với khả năng thực thi phát hiện đối tượng theo thời gian thực.
Module đánh giá hiệu suất tính toán và ghi nhận các chỉ số của mô hình như IoU, ma trận nhầm lẫn, độ chính xác và độ phủ, kèm theo các chức năng trực quan hóa. Module callback tích hợp các công cụ như NeptuneAI để ghi lại dữ liệu thực nghiệm trong quá trình huấn luyện.
Mã nguồn Demo phân đoạn
Phần này trình bày cách triển khai tính năng phân đoạn mà không cần giao diện web, cho phép tích hợp trực tiếp vào các dự án khác. Các chức năng chính bao gồm phân đoạn hình ảnh tĩnh, video và camera theo thời gian thực, cùng với các phép toán hình học như trích xuất contour, phân loại đối tượng, tính chu vi, diện tích, độ tròn và trích xuất màu sắc.
import random
import cv2
import numpy as np
from PIL import ImageFont, ImageDraw, Image
from hashlib import md5
from segmentation_model import BranchSegmenter
from class_labels import CategoryList
def generate_color_by_label(label_name):
"""Sinh màu sắc dựa trên tên nhãn."""
hash_value = int(md5(label_name.encode()).hexdigest(), 16)
return (hash_value % 255, (hash_value >> 8) % 255, (hash_value >> 16) % 255)
def compute_contour_area(contour_points):
"""Tính diện tích đa giác từ các điểm contour."""
return cv2.contourArea(contour_points.astype(np.float32))
def render_chinese_text(image, text_content, position, font_size=20, text_color=(255, 0, 0)):
"""Vẽ văn bản tiếng Trung lên hình ảnh."""
image_array = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
draw_context = ImageDraw.Draw(image_array)
font_object = ImageFont.truetype("simsun.ttc", font_size, encoding="unic")
draw_context.text(position, text_content, font=font_object, fill=text_color)
return cv2.cvtColor(np.array(image_array), cv2.COLOR_RGB2BGR)
def scale_adaptive_parameters(image_dimensions, base_reference=1000):
"""Điều chỉnh tham số động theo kích thước hình ảnh."""
max_dimension = max(image_dimensions)
return max_dimension / base_reference
def visualize_segmentation_results(image, detection_info, overlay_alpha=0.2):
"""Vẽ kết quả phân đoạn lên hình ảnh gốc."""
class_name = detection_info['class_name']
bounding_box = detection_info['bbox']
confidence_score = detection_info['score']
class_index = detection_info['class_id']
segmentation_mask = detection_info['mask']
scale_factor = scale_adaptive_parameters(image.shape[:2])
vertical_spacing = int(20 * scale_factor)
if segmentation_mask is None:
x_left, y_top, x_right, y_bottom = bounding_box
target_region_area = (x_right - x_left) * (y_bottom - y_top)
cv2.rectangle(image, (x_left, y_top), (x_right, y_bottom),
color=(0, 0, 255), thickness=int(3 * scale_factor))
image = render_chinese_text(image, class_name,
(x_left, y_top - int(30 * scale_factor)),
font_size=int(35 * scale_factor))
vertical_offset = int(50 * scale_factor)
else:
mask_coordinates = np.concatenate(segmentation_mask)
target_region_area = compute_contour_area(mask_coordinates)
mask_color = generate_color_by_label(class_name)
try:
blended_layer = image.copy()
cv2.fillPoly(blended_layer, [mask_coordinates.astype(np.int32)], mask_color)
image = cv2.addWeighted(blended_layer, 0.3, image, 0.7, 0)
cv2.drawContours(image, [mask_coordinates.astype(np.int32)], -1,
(0, 0, 255), thickness=int(8 * scale_factor))
# Tính toán các thông số hình học
contour_area = cv2.contourArea(mask_coordinates.astype(np.int32))
contour_perimeter = cv2.arcLength(mask_coordinates.astype(np.int32), True)
# Tính độ tròn
if contour_perimeter > 0:
circularity_measure = (4 * np.pi * contour_area) / (contour_perimeter ** 2)
else:
circularity_measure = 0.0
# Trích xuất màu sắc trung bình
mask_canvas = np.zeros(image.shape[:2], dtype=np.uint8)
cv2.drawContours(mask_canvas, [mask_coordinates.astype(np.int32)], -1, 255, -1)
color_pixels = cv2.findNonZero(mask_canvas)
# Vẽ tên lớp
min_coords = np.min(mask_coordinates, axis=0).astype(int)
image = render_chinese_text(image, class_name,
(min_coords[0], min_coords[1] - int(30 * scale_factor)),
font_size=int(35 * scale_factor))
vertical_offset = int(50 * scale_factor)
# Hiển thị các thông số
metrics_data = [
("Area", contour_area),
("Perimeter", contour_perimeter),
("Circularity", circularity_measure),
("Color", color_str)
]
for metric_index, (metric_label, metric_value) in enumerate(metrics_data):
# Xử lý hiển thị từng chỉ số
pass
except Exception as processing_error:
pass
return image, target_region_area
def execute_frame_processing(model_instance, input_frame):
"""Xử lý từng khung hình đầu vào."""
preprocessed = model_instance.preprocess(input_frame)
prediction = model_instance.predict(preprocessed)
detection = prediction[0] if detection is not None and len(detection)
if detection:
processed_results = model_instance.postprocess(prediction)
for result_data in processed_results:
input_frame, _ = visualize_segmentation_results(input_frame, result_data)
return input_frame
if __name__ == "__main__":
class_categories = CategoryList
segmentation_model = BranchSegmenter()
segmentation_model.load_weights("./weights/yolov8s-seg.pt")
# Xử lý camera realtime
camera_capture = cv2.VideoCapture(0)
while camera_capture.isOpened():
capture_status, frame_data = camera_capture.read()
if not capture_status:
break
# Xử lý từng khung hình
# Xử lý hình ảnh tĩnh
image_file_path = './icon/sample.jpg'
input_image = cv2.imread(image_file_path)
if input_image is not None:
processed_frame = execute_frame_processing(segmentation_model, input_image)
# Lưu hoặc hiển thị kết quả
# Xử lý video
video_file_path = '' # Đường dẫn video đầu vào
video_capture = cv2.VideoCapture(video_file_path)
while video_capture.isOpened():
frame_status, frame_content = video_capture.read()
# Xử lý từng khung hình videoTổng kết
Hệ thống phân đoạn hình ảnh cành cây dựa trên YOLOv8 cải tiến mang lại giải pháp hiệu quả cho các tác vụ nhận diện và phân đoạn trong lĩnh vực nông nghiệp và sinh thái. Với khả năng xử lý đa dạng nguồn dữ liệu và tính toán các thông số hình học chính xác, hệ thống cung cấp nền tảng vững chắc cho việc giám sát và phân tích thực vật tự động.