# Segmentation-Model

**Repository Path**: XiongfeiBai/segmentation-model

## Basic Information

- **Project Name**: Segmentation-Model
- **Description**: using DNN model for defect segmentation
- **Primary Language**: Python
- **License**: Apache-2.0
- **Default Branch**: master
- **Homepage**: https://gitee.com/XiongfeiBai/segmentation-mode
- **GVP Project**: No

## Statistics

- **Stars**: 2
- **Forks**: 0
- **Created**: 2022-09-16
- **Last Updated**: 2024-03-11

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Segmentation-Model

## 介绍
1.若训练数据为工业缺陷数据集，如表面划痕，焊缝缺陷等，且缺陷样本数量较少，正常样本数量较多，

可以采用泊松融合的方法进行数据增强。

2.使用Unet模型对焊缝金相组织缺陷数据进行缺陷分割。

## 软件架构
Pytorch 1.8.0

Windows 10


## 安装教程

git  clone [https://gitee.com/XiongfeiBai/segmentation-model.git](https://gitee.com/XiongfeiBai/segmentation-model)

## 使用说明

### 1.代码结构

```
|--checkpoints             //存储生成的权重文件
|--PossionFusion
|	|--args.py             //参数文件
|   |--get_data.py         //将json文件转变成图片数据
|   |--get_split_data.py   //得到img和相应的mask文件夹
|   |--OTSU.py             //采用OTSU方法进行阈值分割
|   |--possion_fusion.py   //泊松融合脚本
|   |--run.py              //运行该文件，直接会按照相应的方法进行泊松融合
|   |--utils.py
|--data
|   |--data                //储存训练文件
|	|--json_data           //储存json文件
|   |--get_data.py         //json转变成图片数据
|   |--get_label.py        //得到相应mask数据
|   |--data_split.py       //划分train,val,test数据集，并生成相应txt文件
|--dataloader
|   |--dataload.py         //加载数据
|--model
|   |--Unet-model          //Unet模型
|   |--utils
|--loss
|   |-loss.py
|--input                   //储存待预测的图片
|--output                  //储存预测结果
|--args.py
|--predict.py
|--README.md
|--test.py
|--predict.py
```

### 2.泊松融合扩充数据

#### 2.1  缺陷提取

缺陷提取的目的主要是为了得到缺陷区域的二值掩膜（mask）,后续根据该mask 中的白色区域决定参与融合的缺陷区域大小和位置。

以下提供三种方法获取相应的mask:

1:直接使用裁剪出缺陷区域，并以该整个区域作为融合区域（融合效果一般，但操作简单）

		1. 将正常样本放入PossionFusion/normal_data文件夹中
		1. 将缺陷样本放入PossionFusion/defect_data中

![输入图片说明](image-20220922164453020.png)

2:对缺陷区域进行二值分割，得到较为精细的mask。

  		1. 将正常样本放入PossionFusion/normal_data文件夹中
  		2. 将缺陷样本放入PossionFusion/defect_data文件夹中
  		3. 将mask放入PossionFusion/mask文件夹中

![输入图片说明](image-20220922150049642.png)

3.若既想要较为准确的mask,但通过二值分割的方法很难得到较好的分割结果，可参考该工程化方法，具体如下：

  		1. 将正常样本放入PossionFusion/normal_data文件夹中
  		2. 利用labelme对缺陷区域进行标注，将得到的.json文件放入PossionFusion/json_data

#### 2.2  缺陷融合

1.按照上述三种方法的一种放好相应的数据

2.修改PossionFusion/args.py中的配置，根据所选的方法，修改method=1（2或3），fusion-only设置为False

3.运行一次run.py，后将fusion-only设置为True,不断运行run.py即可

4.最终生成的融合图片存放在gen_data文件夹中

### 3. data数据集加载:

```python
data的存储结构如下:
|--data
	|--data
	|--json_data
    |--get_data.py
    |--get_label.py
    |--data_split.py
1.将制作数据集得到的json放入json_data文件夹中.
2.运行get_data.py,会将json文件解析得到的图像放入data/data中.
3.运行get_label.py,将会得到训练所需要的mask数据.
4.运行data_split.py,对数据划分,划分成train,val,test数据集,标号放入对应的txt文件中.
```

### 4.训练数据集

```python
1.对args进行参数设置:
对args.py中的参数进行设置

parser = argparse.ArgumentParser()
parser.add_argument("--data_path", type=str, default="data", help="the directory of data")
parser.add_argument("--split", type=list, default=[0.6, 0.2, 0.2])
parser.add_argument("--epoches", type=int, default=50)
parser.add_argument("--batch_size", type=int, default=4)
parser.add_argument("--num_workers", type=int, default=0)
parser.add_argument("--lr", type=float, default=1e-5)
parser.add_argument("--LR_Schedule", type=str, default="Expoential LR")
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
parser.add_argument("--save_checkpoint", type=str, default="checkpoints",help="the directory of ckpt")
parser.add_argument("--ckpt_path", type=str, default="checkpoints/val_loss_0.4351.pth")
parser.add_argument("--input", type=str, default="input",help="the directory of images needed to be predicted")
parser.add_argument("--output", type=str, default="output",help="the directory of result of prediction")

2.运行train.py

```

### 5.测试数据集

```python
运行test.py
```

### 6.进行预测

```python
1.将待预测的图片放入input文件夹中
2.运行predict.py
```

### 7.结果展示

![输入图片说明](test_img.png)![输入图片说明](img5.png)