2023年亚太杯数学建模A题水果采摘机器人的图像识别功能（基于yolov5的苹果分割）

news2024/10/5 17:16:50

注：.题中附录并没有给出苹果的标签集，所以需要我们自己通过前4问得到训练的标签集，采用的是yolov5 7.0 版本，该版本带分割功能

一：关于数据集的制作：

clc;
close all;
clear;
%-----这个是生成yolov5 数据集的--------
% 图像文件夹路径
folder_path = 'E:/新建文件夹/yatai/Attachment/Apple/';
% 图像文件列表
image_files = dir(fullfile(folder_path, '*.jpg')); % 假设所有图片都是jpg格式

% 解析文件名中的数字，并转换为数值类型
numbers = cellfun(@(x) sscanf(x, '%d.jpg'), {image_files.name});


% 根据解析出的数字对文件列表进行排序
[~, sorted_idx] = sort(numbers);
image_files = image_files(sorted_idx);
% 存储每张图片苹果数量的数组
apple_counts = zeros(length(image_files), 1);

% 存储每张图片的平均成熟度评分
average_red_intensity_ratio_per_image = zeros(length(image_files), 1);




% 确保输出文件夹存在
output_folder = 'E:\新建文件夹\yatai\Attachment\Attachment 2\APPlemasktxt';
if ~exist(output_folder, 'dir')
    mkdir(output_folder);
end

% 存储每张图片的平均苹果质量评分
average_quality_scores_per_image = zeros(length(image_files), 1);
% 遍历每张图片
for i = 1: length(image_files) %2  
    % 读取图像
    img = imread(fullfile(folder_path, image_files(i).name));

     ·······省略了部分代码

    % 给分割的对象标记不同的标签
    labelled_img = bwlabel(binary_img);
%     figure;
%     在原始图像上绘制分割结果
%     imshow(img);
%     hold on;
    colors=['b' 'g' 'r' 'c' 'm' 'y'];
    for k = 1:length(unique(labelled_img)) - 1
        boundary = bwboundaries(labelled_img == k);
        for b = 1:length(boundary)
            plot(boundary{b}(:,2), boundary{b}(:,1), colors(mod(k,length(colors))+1), 'LineWidth', 2);
        end
    end
    % title('Segmented Apples');
    % hold off;

    % 计数分割后的苹果
    number_of_apples = max(labelled_img(:));
    disp(['Number of segmented apples: ', num2str(number_of_apples)]);


    apple_counts(i) = number_of_apples;

    % 打印当前图片的苹果数量
    fprintf('Image %d (%s): %d apples detected.\n', i, image_files(i).name, number_of_apples);

    %下面是制作分割的数据集

    % 给分割的对象标记不同的标签
    labelled_img = bwlabel(binary_img);

    % 准备写入YOLOv5格式的分割轮廓点文件
  
    % 根据图像文件名创建对应的txt文件名
    baseFileName = sprintf('%d.txt', i);
    txt_filename = fullfile(output_folder, baseFileName);
    fileID = fopen(txt_filename, 'w');

    % 确保文件已成功打开
    if fileID == -1
        error('Cannot open file %s for writing.', txt_filename);
    end


    % 获取图像尺寸
    img_height = size(img, 1);
    img_width = size(img, 2);

    % 遍历每个苹果，写入轮廓点信息
    for k = 1:max(labelled_img(:))
       [B,~] = bwboundaries(labelled_img == k, 'noholes');
        contours = B{1}; % 取第一组轮廓点

        % 检查contours的尺寸
        if size(contours, 2) == 2 % 确保contours有两列
            % 转换为归一化坐标
            contours_normalized = contours ./ [img_height,  img_width];

            % 写入文件
            fprintf(fileID, '0 '); % 假设苹果的类别ID为0
            for p = 1:size(contours_normalized, 1)
%                 fprintf('Plotting point at (%f, %f)\n', contours_normalized(p, 2), contours_normalized(p, 1)); % 调试信息
                fprintf(fileID, '%f %f ', contours_normalized(p, 2), contours_normalized(p, 1));
            end
            fprintf(fileID, '\n');
        else
            warning('Contour for apple %d in image %d does not have correct dimensions.', k, i);
        end
    end

    fclose(fileID);
    

end

二：关于yolov5 7.0 的训练：

我的电脑是3080 训练了20轮测试，下面就是部分测试的结果

下面是关于数据集的划分代码

'''
Descripttion: split_img.py
version: 1.0
Author: UniDome
Date: 2022-04-20 16:28:45
LastEditors: UniDome
LastEditTime: 2022-04-20 16:39:56
'''
import os, shutil, random
from tqdm import tqdm


def split_img(img_path, label_path, split_list):
    try:  # 创建数据集文件夹
        Data = 'E:/新建文件夹/yatai/Attachment/Attachment 2/output'
        os.mkdir(Data)

        train_img_dir = Data + '/images/train'
        val_img_dir = Data + '/images/val'
        test_img_dir = Data + '/images/test'

        train_label_dir = Data + '/labels/train'
        val_label_dir = Data + '/labels/val'
        test_label_dir = Data + '/labels/test'

        # 创建文件夹
        os.makedirs(train_img_dir)
        os.makedirs(train_label_dir)
        os.makedirs(val_img_dir)
        os.makedirs(val_label_dir)
        os.makedirs(test_img_dir)
        os.makedirs(test_label_dir)

    except:
        print('文件目录已存在')

    train, val, test = split_list
    all_img = os.listdir(img_path)
    all_img_path = [os.path.join(img_path, img) for img in all_img]
    # all_label = os.listdir(label_path)
    # all_label_path = [os.path.join(label_path, label) for label in all_label]
    train_img = random.sample(all_img_path, int(train * len(all_img_path)))
    train_img_copy = [os.path.join(train_img_dir, img.split('\\')[-1]) for img in train_img]
    train_label = [toLabelPath(img, label_path) for img in train_img]
    train_label_copy = [os.path.join(train_label_dir, label.split('\\')[-1]) for label in train_label]
    for i in tqdm(range(len(train_img)), desc='train ', ncols=80, unit='img'):
        _copy(train_img[i], train_img_dir)
        _copy(train_label[i], train_label_dir)
        all_img_path.remove(train_img[i])
    val_img = random.sample(all_img_path, int(val / (val + test) * len(all_img_path)))
    val_label = [toLabelPath(img, label_path) for img in val_img]
    for i in tqdm(range(len(val_img)), desc='val ', ncols=80, unit='img'):
        _copy(val_img[i], val_img_dir)
        _copy(val_label[i], val_label_dir)
        all_img_path.remove(val_img[i])
    test_img = all_img_path
    test_label = [toLabelPath(img, label_path) for img in test_img]
    for i in tqdm(range(len(test_img)), desc='test ', ncols=80, unit='img'):
        _copy(test_img[i], test_img_dir)
        _copy(test_label[i], test_label_dir)


def _copy(from_path, to_path):
    shutil.copy(from_path, to_path)


def toLabelPath(img_path, label_path):
    img = img_path.split('\\')[-1]
    label = img.split('.jpg')[0] + '.txt'
    return os.path.join(label_path, label)


def main():
    img_path = r'E:\新建文件夹\yatai\Attachment\Apple'
    label_path = r'E:\新建文件夹\yatai\Attachment\Attachment 2\APPlemasktxt'
    split_list = [0.7, 0.2, 0.1]  # 数据集划分比例[train:val:test]
    split_img(img_path, label_path, split_list)


if __name__ == '__main__':
    main()