unet module

.gitignore

 img/*
 out/*
 __pycache__/
-.idea/*
\ No newline at end of file
+.idea/*
+data/imgs/*
+data/masks/*
+data/output/*
\ No newline at end of file

README.md

@@ -5,30 +5,70 @@
 ## 目录结构
 ```.
 ├── README.md
-├── img
-│   └── folders
-│       └── imgs
-├── out
-│   ├── csv
-│   │   └── result
-│   └── graph
-│       └── processed images
-├── main.py
-└── util.py
+├── data
+│   ├── module
+│   │   └── trained modules
+│   ├── imgs
+│   │   └── origin images
+│   ├── masks
+│   │   └── ground truth masks value (0~1)
+│   └── output
+│       ├── csv
+│       │   └── result
+│       └── graph
+│          └── processed images
+└── other codes
 ```
-## 参数
+## 主程序入口
+### 参数
 * -m ,--method  : 0 使用Kmeans，1 使用阈值法(butterworth滤波),2 使用阈值法（fft）。默认Kmeans
 * -c ,--core    : Kmeans分为几类，默认5，仅对Kmeans法有效
 * -p ,--process : 使用线程数量，默认8，仅对阈值法有效
 
-## 示例
+### 示例
 
 - `python main.py -m 1 -p 16`
 - `python main.py -m 0 -c 8`
 
+
+## 使用Unet模型
+### 训练
+```shell script
+> python train.py -h
+usage: train.py [-h] [-e E] [-b [B]] [-l [LR]] [-f LOAD] [-s SCALE] [-v VAL]
+
+Train the UNet on images and target masks
+
+optional arguments:
+  -h, --help            show this help message and exit
+  -e E, --epochs E      Number of epochs (default: 5)
+  -b [B], --batch-size [B]
+                        Batch size (default: 1)
+  -l [LR], --learning-rate [LR]
+                        Learning rate (default: 0.1)
+  -f LOAD, --load LOAD  Load model from a .pth file (default: False)
+  -s SCALE, --scale SCALE
+                        Downscaling factor of the images (default: 0.5)
+  -v VAL, --validation VAL
+                        Percent of the data that is used as validation (0-100)
+                        (default: 15.0)
+
+```
+
+### 监控
+使用tensorboard可视化监控
+`tensorboard --logdir=runs`
+
+## 数据
+输入图像存放于imgs文件夹下，单通道灰度图像，8bit，200*200
+标注的mask存放于masks文件夹下，像素0为背景，1为目标
+
+## cli工具
+部分可能用得到的cli工具，主要为python或shell脚本。可根据需要自行修改。
 ## 版本
 
 |版本号|功能|
 |:---:|:---|
 |v0.1|实现了Kmeans、阈值法两种分割方法，提供命令行工具|
-|v0.2|1.更新了滤波算法<br>2.修改了一些小bug<br>3.对部分函数进行解耦|
\ No newline at end of file
+|v0.2|1.更新了滤波算法<br>2.修改了一些小bug<br>3.对部分函数进行解耦|
+|v0.3|1.训练了Unet模型<br>2.增加了部分命令行工具<br>4.重新设计目录结构|
\ No newline at end of file

main.py

@@ -33,8 +33,6 @@ def get_args():
                                      formatter_class = argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument('-m', '--method', metavar = 'M', type = int, default = 0,
                         help = '0 for KMeans; 1 for Threshold; 2 for newThreshold', dest = 'method')
-    #parser.add_argument('-o', '--out', metavar = 'O', type = str, default = './out', help = 'Path to output',
-    #                    dest = 'out')
     parser.add_argument('-c', '--core', metavar = 'C', type = int, default = 5, help = 'Num of cluster', dest = 'core')
     parser.add_argument('-p', '--process', metavar = 'P', type = int, default = 8, help = 'Num of process',
                         dest = 'process')

submit.py

-""" Submit code specific to the kaggle challenge"""
-
-import os
-
-import torch
-from PIL import Image
-import numpy as np
-
-from predict import predict_img
-from unet import UNet
-
-# credits to https://stackoverflow.com/users/6076729/manuel-lagunas
-def rle_encode(mask_image):
-    pixels = mask_image.flatten()
-    # We avoid issues with '1' at the start or end (at the corners of
-    # the original image) by setting those pixels to '0' explicitly.
-    # We do not expect these to be non-zero for an accurate mask,
-    # so this should not harm the score.
-    pixels[0] = 0
-    pixels[-1] = 0
-    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2
-    runs[1::2] = runs[1::2] - runs[:-1:2]
-    return runs
-
-
-def submit(net, gpu=False):
-    """Used for Kaggle submission: predicts and encode all test images"""
-    dir = 'data/test/'
-
-    N = len(list(os.listdir(dir)))
-    with open('SUBMISSION.csv', 'a') as f:
-        f.write('img,rle_mask\n')
-        for index, i in enumerate(os.listdir(dir)):
-            print('{}/{}'.format(index, N))
-
-            img = Image.open(dir + i)
-
-            mask = predict_img(net, img, gpu)
-            enc = rle_encode(mask)
-            f.write('{},{}\n'.format(i, ' '.join(map(str, enc))))
-
-
-if __name__ == '__main__':
-    net = UNet(3, 1).cuda()
-    net.load_state_dict(torch.load('MODEL.pth'))
-    submit(net, True)