# pytorch-rubbish-classification

**Repository Path**: fhzwt/pytorch-rubbish-classification

## Basic Information

- **Project Name**: pytorch-rubbish-classification
- **Description**: 面向垃圾分类回收的图像分类任务实验
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 1
- **Forks**: 1
- **Created**: 2022-11-29
- **Last Updated**: 2023-05-18

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

## 说明

#### 0、依赖

```
pytorch==1.12.1
scikit-image==0.19.3
numpy==1.23.4
torchvision==0.13.1
matplotlib==3.5.2
mxnet-cu101
PyWavelets==1.0.0
opecv-python==4.6.0
sklearn==0.19.3
pillow==9.2.0
# 可选
tensorboard==2.6.0
```

#### 1、准备数据集

解压RubbishClassification.zip，删除其中后缀为.json .py文件。

修改rubbish_dataset.py中的dataset_split()函数中的数据集路径。以下代码仅供参考。请根据实际情况修改。

```python
# 原数据集位置
origin_path = r'.\RubbishClassification' 
# 训练集位置，不存在会自动创建
train_path = r'.\train'
# 测试集位置，不存在会自动创建
test_path = r'.\test'
# 测试集位置，不存在会自动创建
valid_path = r'.\valid'
```

get_training_dataloader_rubbish函数中的path修改为与train_path相同

get_test_dataloader_rubbish函数中的path修改为与test_path相同

get_valid_dataloader_rubbish函数中的path修改为与valid_path相同

**运行rubbish_dataset.py，划分数据集、测试集、验证集。**

#### 2、运行tensorbard（可选）

安装好tensorbard，在终端运行

```
$ tensorboard --logdir='runs' --port=6006 --host='localhost'
```

#### 3、训练神经网络模型

在终端运行

```
# use gpu to train vgg16
$ python train.py -net vgg16 -gpu
```

网络参数可替换为

```
squeezenet
mobilenetv2
shufflenetv2
vgg16
googlenet
resnet18
densenet121
```

#### 4、测试模型

在终端运行

```
$ python test.py -net vgg16 -weights path
```

path为checkpoint对应文件夹下.pth文件路径

#### 5、微调网络

在终端运行

```
$ python fine_tune.py -net vgg16 -gpu
```

网络参数可替换为

```
vgg16
googlenet
resnet18
```

#### 6、测试微调网络

在终端运行

```
$ python test_ft.py -net vgg16 -weights path
```

path为checkpoint对应文件夹下.pth文件路径

#### 7、svm

修改文件中的path为未划分的数据集地址

运行即可

#### 备注

1.最终模型保存在checkpoint目录下，训练过程保存在run目录下。

2.如果需要修改分辨率或者使用vgg18微调，需要调整部分代码。

#### 参考
https://github.com/weiaicunzai/pytorch-cifar100

https://github.com/creafz/pytorch-cnn-finetune