鱼的个体识别-数据集处理

目标

对数据集进行处理，使用SAM模型自动将鱼抠出。

图片预处理

将图片全部预处理为640*640大小。

分割方法

方法介绍

使用MobileSAM模型，根据12个点位进行分割。

导入包

1
2
3
4
5
import os
from glob import glob
import cv2
import numpy as np
from ultralytics import SAM

设置文件路径

文件目录结构为

1
2
3
4
5
6
- fishImages
  - 1
    - *.JPG
    ...
  - 2
  - ...

读取所有图片路径

1
2
3
4
imagesPath = '.\\fishImages'
imagesPathList = glob(imagesPath+'\\**\\'+ '*.JPG', recursive=True)

print(imagesPathList)

输出结果为

1
['.\\fishImages\\1\\G0285619.JPG', ... ]

加载模型

模型使用的 MobileSAM 模型， 当前使用设备为轻薄本，只有一个2G显卡，显存不够，只能使用cpu进行分割。

1
2
model = SAM("mobile_sam.pt")
model.to("cpu")

图片保存方法

由于根据12个点位进行的分割，因此将点位标号加入到原始文件名末尾。

路径结构为：

1
2
3
4
5
6
- datasets
    - 1
        *1.png
        ...
    - 2
    ...
1
2
3
4
5
6
7
8
9
datasetsPath = '.\\datasets'
def save_img(imagesPath, i, masked_image):
    num = imagesPath.split('\\')[-2]
    imageName = imagesPath.split('\\')[-1]
    savePath = datasetsPath + '\\' + num
    if(not os.path.exists(savePath)):
        os.mkdir(savePath) 
    save_img_name = savePath + '\\' + imageName.split('.')[-2] + str(i) + '.' + 'png'
    cv2.imwrite(save_img_name, masked_image)

分割方法

读取图片路径，根据点位和标签进行分割。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def cut_out_img(imagePath, labels, points):
    orig_img = cv2.imread(imagePath)
    result = model.predict(source=imagePath, 
                           labels=labels,
                           points=points)
    i = 0
    for d in result[0].masks.data:
        # 得到掩码
        mask = d.cpu().numpy()
        masked_image = np.zeros_like(orig_img)
        masked_image[mask] = orig_img[mask]
        # 设置透明背景
        alpha_channel = np.ones(mask.shape, dtype=orig_img.dtype) * 255
        masked_image_rgba = np.dstack((masked_image, alpha_channel)) 
        masked_image_rgba[~mask] = (0, 0, 0, 0)
        # 保存图片
        save_img(imagePath, i, masked_image_rgba)
        i = i+1

点位设置

设置了12个点位，labels数值中的值为0或1， 0为背景， 1为前景。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
points = np.array([
    [300, 470],
    [540, 550],
    [260, 300],
    [300, 380],
    [300, 340],
    [280, 210],
    [350, 430],
    [280, 250],
    [450, 320],
    [340, 470],
    [340, 250],
    [390, 430]
])

labels = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

进行分割

1
2
for imagePath in imagesPathList:
    cut_out_img(imagePath, labels, points)