2023年亚太杯APMCM数学建模大赛A题水果采摘机器人的图像识别

2023年亚太杯APMCM数学建模大赛

A题 水果采摘机器人的图像识别

原题再现

  中国是世界上最大的苹果生产国，年产量约3500万吨。同时，中国也是世界上最大的苹果出口国，世界上每两个苹果中就有一个是中国出口的，世界上超过六分之一的苹果是中国出口的。中国提出了“一带一路”倡议，这是建设未来共享的国际社会的重要支柱。由于这一举措，越南、孟加拉国、菲律宾、印度尼西亚等沿线国家已成为中国苹果的主要出口目的地。

  苹果采摘主要依靠手工采摘。苹果成熟后，几天内就需要大量的采摘工人。但大多数当地农民在自己的果园里种植苹果。此外，农业工人的老龄化和年轻人离开村庄工作的现象导致了苹果采摘季节的劳动力短缺。为了解决这一问题，中国从2011年左右开始研究能够摘苹果的机器人，并取得了重大进展。
  然而，由于果园环境与受控实验环境的差异，各种苹果采摘机器人在世界范围内的推广应用还不够理想。在复杂、非结构化的果园环境中，现有的机器人大多无法准确识别障碍物，如“树叶遮挡”、“树枝遮挡”、“水果遮挡”、“混合遮挡”等，如果直接采摘苹果而不根据实际场景做出准确判断，则损坏水果的风险很高，甚至对采摘手和机械臂造成伤害。这对收获效率和果实质量产生不利影响，导致更大的损失。此外，对不同采收果实的识别和分类也非常重要，如分类、加工、包装、运输等程序。然而，许多水果的颜色、形状和大小与苹果非常相似，这给苹果的采后识别带来了很大的困难。
  本次竞赛的目的是通过分析和提取标记水果图像的特征，建立一个识别率高、速度快、准确度高的苹果图像识别模型，并对图像进行数据分析，如自动计算图像中苹果的数量、位置、成熟度，估计图像中苹果的质量等。具体任务如下：

  问题1：数苹果
  基于附件1中提供的已采收苹果图像数据集，提取图像特征，建立数学模型，统计每个图像中的苹果数，绘制附件1中所有苹果分布的直方图。

  问题2：估计苹果的位置
  根据附件1中提供的已收获苹果的图像数据集，以图像左下角为坐标原点，识别每个图像中苹果的位置，并绘制附件1中所有苹果几何坐标的二维散点图。

  问题3：估计苹果的成熟状态
  根据附件1提供的已采收苹果图像数据集，建立数学模型，计算每幅图像中苹果的成熟度，并绘制附件1中所有苹果成熟度分布的直方图。

  问题4：估计苹果的质量
  根据附件1中提供的已收获苹果的图像数据集，以图像左下角为坐标原点，计算每幅图像中苹果的二维面积，估计苹果的质量，并绘制附件1中所有苹果的质量分布直方图。

  问题5：苹果的识别
  基于附件2提供的采集水果图像数据集，提取图像特征，训练苹果识别模型，识别附件3中的苹果，并绘制附件3中所有苹果图像ID号的分布直方图。

import cv2
import torch
import torchvision
import numpy as np
from torchvision import transforms, utils
from matplotlib import pyplot as plt

CLASS_NAMES = ['background', 'apple']
COLORS = np.random.uniform(0, 255, size=(len(CLASS_NAMES), 3))

def load_model():
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    model.eval()
    return model
    
def load_and_process_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    transformed = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((800, 800)),
    ])
    image = transformed(image)
    return image.unsqueeze(0)



model = load_model()

import glob
import os
def visualize(image, boxes, labels, title,i):
    image = image.permute(1, 2, 0).numpy()
    fig = plt.figure(figsize=(8, 8))
    plt.imshow(image)
    ax = plt.gca()
    for i in range(len(boxes)):
        xmin, ymin, xmax, ymax = boxes[i]
        color = 'r'
        rect = plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, edgecolor=color, linewidth=2)
        ax.add_patch(rect)
        text = '{:.2f}'.format(scores[i])
        plt.text(xmin, ymin-5, text, fontsize=8, bbox=dict(facecolor=color, alpha=0.5, pad=1), color='white')
    plt.axis('off')
    plt.title(title)
    plt.savefig('img/Q1_{}.png'.format(i),dpi=100)
    plt.show()

# 循环读取图片
image_paths = sorted(glob.glob(os.path.join('Attachment/Attachment 1', '*.jpg')))
i=0
test = []
count_list = []
for image_path in image_paths:
    image = load_and_process_image(image_path)
    det_pred = model(image)
    boxes = det_pred[0]['boxes'].detach().numpy()
    scores = det_pred[0]['scores'].detach().numpy()
    labels = det_pred[0]['labels'].detach().numpy()
    
    threshold = 0.5
    idx = scores > threshold
    boxes = boxes[idx]
    scores = scores[idx]
    labels = labels[idx]
    
    num_apples = len(boxes)
    count_list.append(num_apples)
    print('第{}张图片中苹果的数量为{}个'.format(i, num_apples))
    apple_locs = []
    for j in range(num_apples):
        xmin, ymin, xmax, ymax = boxes[j]
        apple_loc = (xmin, ymin, xmax, ymax)
        apple_locs.append(apple_loc)
    if i <= 3:
        visualize(image[0], boxes, labels, 'Image {}'.format(i),i)
    i+=1
# 绘制直方图
plt.hist(count_list, bins=range(max(count_list) + 2))
plt.xlabel("Apple count")
plt.ylabel("Image count")
plt.savefig('apple_number.png',dpi=300)
plt.show()