diff --git a/README.md b/README.md
index beead7995..ce44cb214 100644
--- a/README.md
+++ b/README.md
@@ -29,11 +29,7 @@ python iann/__main__.py
 - [x] 解决图片放大后缩小，顶部定格显示不能居中的问题
 - [x] 解决放大图像后平移速度变快
 - [x] 解决图像切换后窗口缩放没有清除（顺便缩放到显示框最佳）
-
-## 新提交
-
-1. 模型未加载时直接打开文件夹会报错，发现在打开文件夹直接加载模型前就访问了controller的属性，而模型加载时才创建的controller；这里改了一下代码的顺序。
-2. 直接加载模型会报错，发现在加载模型时默认打开的图像，当未选择时图像为None；这里加了一个判断，如果直接加载模型，当图像为None就不加载。
-3. 每次加载图像时根据图像大小重新设定scene的区域，这样缩小图像和切换不会造成不能居中的问题。
-4. 添加最佳缩放，每次切换图片后，自适应缩放到最佳显示。
-5. update_image时判断是否应该缩放重置，避免在点击点时也进行了图像缩放重置。
+- [ ] 训练代码整理和训练配置的抽出（初步实现）
+- [ ] 重新审查修改按键图标，确保图标完整清晰
+- [ ] 界面配色以paddle蓝色为主，同时做一个ico窗口图标
+- [ ] 标签可修改？
\ No newline at end of file
diff --git a/iann/data/mdiy.py b/iann/data/mdiy.py
new file mode 100644
index 000000000..0302b48f7
--- /dev/null
+++ b/iann/data/mdiy.py
@@ -0,0 +1,44 @@
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+from .base import ISDataset
+
+
+class MyDataset(ISDataset):
+    def __init__(self, dataset_path, folder_name,
+                 images_dir_name, masks_dir_name,
+                 **kwargs):
+        super(MyDataset, self).__init__(**kwargs)
+
+        self.dataset_path = Path(dataset_path) / folder_name
+        self._images_path = self.dataset_path / images_dir_name
+        self._insts_path = self.dataset_path / masks_dir_name
+
+        self.dataset_samples = [x.name for x in sorted(self._images_path.glob('*.*'))]
+        self._masks_paths = {x.stem: x for x in self._insts_path.glob('*.*')}
+
+    def get_sample(self, index):
+        image_name = self.dataset_samples[index]
+        image_path = str(self._images_path / image_name)
+        mask_path = str(self._masks_paths[image_name.split('.')[0]])
+
+        image = cv2.imread(image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        instances_mask = np.max(cv2.imread(mask_path).astype(np.int32), axis=2)
+        instances_mask[instances_mask > 0] = 1
+
+        instances_ids = [1]
+
+        instances_info = {
+            x: {'ignore': False}
+            for x in instances_ids
+        }
+
+        return {
+            'image': image,
+            'instances_mask': instances_mask,
+            'instances_info': instances_info,
+            'image_id': index
+        }
\ No newline at end of file
diff --git a/iann/model/loss.py b/iann/model/loss.py
new file mode 100644
index 000000000..ba7315255
--- /dev/null
+++ b/iann/model/loss.py
@@ -0,0 +1,133 @@
+import numpy as np
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+import util.util as U
+
+
+class NormalizedFocalLossSigmoid(nn.Layer):
+    def __init__(self, axis=-1, alpha=0.25, gamma=2,
+                 from_logits=False, batch_axis=0,
+                 weight=None, size_average=True, detach_delimeter=True,
+                 eps=1e-12, scale=1.0,
+                 ignore_label=-1):
+        super(NormalizedFocalLossSigmoid, self).__init__()
+        self._axis = axis
+        self._alpha = alpha
+        self._gamma = gamma
+        self._ignore_label = ignore_label
+        self._weight = weight if weight is not None else 1.0
+        self._batch_axis = batch_axis
+
+        self._scale = scale
+        self._from_logits = from_logits
+        self._eps = eps
+        self._size_average = size_average
+        self._detach_delimeter = detach_delimeter
+        self._k_sum = 0
+
+    def forward(self, pred, label, sample_weight=None):
+        one_hot = label > 0
+        sample_weight = label != self._ignore_label
+        sample_weight = sample_weight.astype('float32')
+
+        if not self._from_logits:
+            pred = F.sigmoid(pred)
+
+        x = sample_weight * 0.5
+        y = (1 - self._alpha) * sample_weight
+        alpha = paddle.where(one_hot, x, y)
+        pt = paddle.where(one_hot, pred, 1 - pred)
+        sample_weight = sample_weight.astype('bool')
+        pt = paddle.where(sample_weight, pt, paddle.ones_like(pt))
+        beta = (1 - pt) ** self._gamma
+        sample_weight = sample_weight.astype('float32')
+        sw_sum = paddle.sum(sample_weight, axis=(-2, -1), keepdim=True)
+        beta_sum = paddle.sum(beta, axis=(-2, -1), keepdim=True)
+        mult = sw_sum / (beta_sum + self._eps)
+
+        if self._detach_delimeter:
+            mult = mult.detach()
+
+        beta = beta * mult
+        ignore_area = paddle.sum((label == self._ignore_label).astype('float32'), axis=tuple(range(1, len(label.shape)))).numpy()
+        sample_mult = paddle.mean(mult, axis=tuple(range(1, len(mult.shape)))).numpy()
+        if np.any(ignore_area == 0):
+            self._k_sum = 0.9 * self._k_sum + 0.1 * sample_mult[ignore_area == 0].mean()
+        loss = -alpha * beta * paddle.log(paddle.mean(pt + self._eps))
+        loss = self._weight * (loss * sample_weight)
+
+        if self._size_average:
+            bsum = paddle.sum(sample_weight, axis=U.get_dims_with_exclusion(len(sample_weight.shape), self._batch_axis))
+            loss = paddle.sum(loss, axis=U.get_dims_with_exclusion(len(loss.shape), self._batch_axis)) / (
+                    bsum + self._eps)
+        else:
+            loss = paddle.sum(loss, axis=U.get_dims_with_exclusion(len(loss.shape), self._batch_axis))
+
+        return self._scale * loss
+
+
+class FocalLoss(nn.Layer):
+    def __init__(self, axis=-1, alpha=0.25, gamma=2,
+                 from_logits=False, batch_axis=0,
+                 weight=None, num_class=None,
+                 eps=1e-9, size_average=True, scale=1.0):
+        super(FocalLoss, self).__init__()
+        self._axis = axis
+        self._alpha = alpha
+        self._gamma = gamma
+        self._weight = weight if weight is not None else 1.0
+        self._batch_axis = batch_axis
+
+        self._scale = scale
+        self._num_class = num_class
+        self._from_logits = from_logits
+        self._eps = eps
+        self._size_average = size_average
+
+    def forward(self, pred, label, sample_weight=None):
+        if not self._from_logits:
+            pred = F.sigmoid(pred)
+
+        one_hot = label > 0
+        pt = paddle.where(one_hot, pred, 1 - pred)
+        t = label != -1
+        alpha = paddle.where(one_hot, self._alpha * t, (1 - self._alpha) * t)
+        beta = (1 - pt) ** self._gamma
+
+        loss = -alpha * beta * paddle.log(paddle.min(pt + self._eps, paddle.ones(1, dtype='float32')))
+        sample_weight = label != -1
+
+        loss = self._weight * (loss * sample_weight)
+
+        if self._size_average:
+            tsum = paddle.sum(label == 1, axis=U.get_dims_with_exclusion(len(label.shape), self._batch_axis))
+            loss = paddle.sum(loss, axis=U.get_dims_with_exclusion(len(loss.shape), self._batch_axis)) / (
+                    tsum + self._eps)
+        else:
+            loss = paddle.sum(loss, axis=U.get_dims_with_exclusion(len(loss.shape), self._batch_axis))
+
+        return self._scale * loss
+
+
+class SigmoidBinaryCrossEntropyLoss(nn.Layer):
+    def __init__(self, from_sigmoid=False, weight=None, batch_axis=0, ignore_label=-1):
+        super(SigmoidBinaryCrossEntropyLoss, self).__init__()
+        self._from_sigmoid = from_sigmoid
+        self._ignore_label = ignore_label
+        self._weight = weight if weight is not None else 1.0
+        self._batch_axis = batch_axis
+
+    def forward(self, pred, label):
+        label = label.reshape(pred.shape)
+        sample_weight = label != self._ignore_label
+        label = paddle.where(sample_weight, label, paddle.zeros_like(label))
+
+        if not self._from_sigmoid:
+            loss = F.relu(pred) - pred * label + F.softplus(-paddle.abs(pred))
+        else:
+            eps = 1e-12
+            loss = -(paddle.log(pred + eps) * label + paddle.log(1. - pred + eps) * (1. - label))
+        loss = self._weight * (loss * sample_weight)
+        return paddle.mean(loss, axis=U.get_dims_with_exclusion(len(loss.shape), self._batch_axis))
+
diff --git a/iann/train/README.md b/iann/train/README.md
new file mode 100644
index 000000000..987db7085
--- /dev/null
+++ b/iann/train/README.md
@@ -0,0 +1,81 @@
+# 训练iann可用的自定义模型
+
+目前已经可以通过简单的配置完成模型训练了，但其中有些设置还不能通过配置文件进行修改。
+
+## 一、数据组织
+
+在需要训练自己的数据集时，目前需要将数据集构造为如下格式，直接放在datasets文件夹中。文件名可以根据要求来进行设置，只需要在配置文件中设定好即可，图像和标签与平时使用的分割图像的用法相同。
+
+```
+datasets
+    |
+    ├── train_data
+    |       ├── img
+    |       |    └── filename_1.jpg
+    |       └── gt
+    |            └── filename_1.png
+    |
+    └── eval_data
+            ├── img
+            |    └── filename_1.jpg
+            └── gt
+                 └── filename_1.png
+```
+
+## 二、训练
+
+直接运行ritm_train.py即可开始训练。
+
+```python
+%cd train
+! python ritm_train.py --config train_config.yaml
+```
+
+目前一些简单的参数已经可以在yaml配置文件中进行自定义设置，不过现阶段仍然不够灵活，可能出现各种问题。
+
+```
+iters: 100000  # 训练轮数
+batch_size: 16  # bs大小
+save_interval: 1000  # 保存间隔
+log_iters: 10  # 打印log的间隔
+worker: 4  # 子进程数
+save_dir: model_output  # 保存路径
+use_vdl: False  # 是否使用vdl
+
+dataset:
+  dataset_path: iann/train/datasets  # 数据集所在路径
+  image_name: img  # 图像文件夹的名称
+  label_name: gt  # 标签文件夹的名称
+
+train_dataset:  # 训练数据
+  crop_size: [320, 480]  # 裁剪大小
+  folder_name: train_data  # 训练数据文件夹的名称
+
+val_dataset:  # 验证数据
+  folder_name: val_data  # 验证数据文件夹的名称
+
+optimizer:
+  type: adam  # 优化器，目前仅可以选择‘adam’和‘sgd’
+
+learning_rate:
+  value_1: 5e-5  # 需要设置两个学习率
+  value_2: 5e-6
+  decay:
+    type: poly  # 学习率衰减，目前仅支持‘poly’，可以修改下面的参数
+    steps: 1000
+    power: 0.9
+    end_lr: 0.0
+
+model:
+  type: deeplab  # 模型名称，目前支持‘hrnet’、‘deeplab’以及‘shufflenet’
+  backbone: resnet18  # 下面的参数是模型对应的参数，可在源码中查看
+  is_ritm: True
+  weights: None  # 加载权重的路径
+```
+
+
+
+### * 说明
+
+1. 这里有个坑，数据不能有没有标签的纯背景，这样找不到正样点训练就会卡住，并且还不报错。
+
diff --git a/iann/train/datasets/train_data/gt/frame_2665.png b/iann/train/datasets/train_data/gt/frame_2665.png
new file mode 100644
index 000000000..0a7b86f40
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2665.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2666.png b/iann/train/datasets/train_data/gt/frame_2666.png
new file mode 100644
index 000000000..b01b63d5e
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2666.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2667.png b/iann/train/datasets/train_data/gt/frame_2667.png
new file mode 100644
index 000000000..c27d07d28
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2667.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2668.png b/iann/train/datasets/train_data/gt/frame_2668.png
new file mode 100644
index 000000000..46f21ccf0
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2668.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2669.png b/iann/train/datasets/train_data/gt/frame_2669.png
new file mode 100644
index 000000000..fe702409a
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2669.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2670.png b/iann/train/datasets/train_data/gt/frame_2670.png
new file mode 100644
index 000000000..c4ab51628
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2670.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2671.png b/iann/train/datasets/train_data/gt/frame_2671.png
new file mode 100644
index 000000000..26f999dda
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2671.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2672.png b/iann/train/datasets/train_data/gt/frame_2672.png
new file mode 100644
index 000000000..60ac08a12
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2672.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2673.png b/iann/train/datasets/train_data/gt/frame_2673.png
new file mode 100644
index 000000000..98423b4a0
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2673.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2674.png b/iann/train/datasets/train_data/gt/frame_2674.png
new file mode 100644
index 000000000..282f12670
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2674.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2675.png b/iann/train/datasets/train_data/gt/frame_2675.png
new file mode 100644
index 000000000..e70988b9a
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2675.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2676.png b/iann/train/datasets/train_data/gt/frame_2676.png
new file mode 100644
index 000000000..e73ca251f
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2676.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2678.png b/iann/train/datasets/train_data/gt/frame_2678.png
new file mode 100644
index 000000000..ce31f4edf
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2678.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2679.png b/iann/train/datasets/train_data/gt/frame_2679.png
new file mode 100644
index 000000000..87444d584
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2679.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2680.png b/iann/train/datasets/train_data/gt/frame_2680.png
new file mode 100644
index 000000000..3a3b12b68
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2680.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2681.png b/iann/train/datasets/train_data/gt/frame_2681.png
new file mode 100644
index 000000000..c47413d7d
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2681.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2682.png b/iann/train/datasets/train_data/gt/frame_2682.png
new file mode 100644
index 000000000..c48cab278
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2682.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2683.png b/iann/train/datasets/train_data/gt/frame_2683.png
new file mode 100644
index 000000000..f1a3861e0
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2683.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2685.png b/iann/train/datasets/train_data/gt/frame_2685.png
new file mode 100644
index 000000000..cfdb280ff
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2685.png differ
diff --git a/iann/train/datasets/train_data/gt/frame_2686.png b/iann/train/datasets/train_data/gt/frame_2686.png
new file mode 100644
index 000000000..3e9b198fb
Binary files /dev/null and b/iann/train/datasets/train_data/gt/frame_2686.png differ
diff --git a/iann/train/datasets/train_data/img/frame_2665.jpg b/iann/train/datasets/train_data/img/frame_2665.jpg
new file mode 100644
index 000000000..832a4da54
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2665.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2666.jpg b/iann/train/datasets/train_data/img/frame_2666.jpg
new file mode 100644
index 000000000..267e6439f
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2666.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2667.jpg b/iann/train/datasets/train_data/img/frame_2667.jpg
new file mode 100644
index 000000000..d0c472e18
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2667.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2668.jpg b/iann/train/datasets/train_data/img/frame_2668.jpg
new file mode 100644
index 000000000..01d9431e0
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2668.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2669.jpg b/iann/train/datasets/train_data/img/frame_2669.jpg
new file mode 100644
index 000000000..3e9004e37
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2669.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2670.jpg b/iann/train/datasets/train_data/img/frame_2670.jpg
new file mode 100644
index 000000000..c1affedbf
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2670.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2671.jpg b/iann/train/datasets/train_data/img/frame_2671.jpg
new file mode 100644
index 000000000..5ae0be220
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2671.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2672.jpg b/iann/train/datasets/train_data/img/frame_2672.jpg
new file mode 100644
index 000000000..7fc4abe16
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2672.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2673.jpg b/iann/train/datasets/train_data/img/frame_2673.jpg
new file mode 100644
index 000000000..a147ad912
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2673.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2674.jpg b/iann/train/datasets/train_data/img/frame_2674.jpg
new file mode 100644
index 000000000..75ceefb9c
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2674.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2675.jpg b/iann/train/datasets/train_data/img/frame_2675.jpg
new file mode 100644
index 000000000..894ae0b3a
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2675.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2676.jpg b/iann/train/datasets/train_data/img/frame_2676.jpg
new file mode 100644
index 000000000..03801522d
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2676.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2678.jpg b/iann/train/datasets/train_data/img/frame_2678.jpg
new file mode 100644
index 000000000..eabb00374
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2678.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2679.jpg b/iann/train/datasets/train_data/img/frame_2679.jpg
new file mode 100644
index 000000000..5d9058446
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2679.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2680.jpg b/iann/train/datasets/train_data/img/frame_2680.jpg
new file mode 100644
index 000000000..ada66278e
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2680.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2681.jpg b/iann/train/datasets/train_data/img/frame_2681.jpg
new file mode 100644
index 000000000..a35bbfb61
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2681.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2682.jpg b/iann/train/datasets/train_data/img/frame_2682.jpg
new file mode 100644
index 000000000..94b6b4892
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2682.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2683.jpg b/iann/train/datasets/train_data/img/frame_2683.jpg
new file mode 100644
index 000000000..2d360d8a1
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2683.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2685.jpg b/iann/train/datasets/train_data/img/frame_2685.jpg
new file mode 100644
index 000000000..ba84c508c
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2685.jpg differ
diff --git a/iann/train/datasets/train_data/img/frame_2686.jpg b/iann/train/datasets/train_data/img/frame_2686.jpg
new file mode 100644
index 000000000..6b5501550
Binary files /dev/null and b/iann/train/datasets/train_data/img/frame_2686.jpg differ
diff --git a/iann/train/datasets/val_data/gt/frame_2687.png b/iann/train/datasets/val_data/gt/frame_2687.png
new file mode 100644
index 000000000..1824bf581
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2687.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2689.png b/iann/train/datasets/val_data/gt/frame_2689.png
new file mode 100644
index 000000000..43c6f7922
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2689.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2690.png b/iann/train/datasets/val_data/gt/frame_2690.png
new file mode 100644
index 000000000..9b2a7686d
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2690.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2691.png b/iann/train/datasets/val_data/gt/frame_2691.png
new file mode 100644
index 000000000..1859aded9
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2691.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2692.png b/iann/train/datasets/val_data/gt/frame_2692.png
new file mode 100644
index 000000000..2ff3eeb7a
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2692.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2693.png b/iann/train/datasets/val_data/gt/frame_2693.png
new file mode 100644
index 000000000..142a40f55
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2693.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2695.png b/iann/train/datasets/val_data/gt/frame_2695.png
new file mode 100644
index 000000000..5bbf13d4a
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2695.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2697.png b/iann/train/datasets/val_data/gt/frame_2697.png
new file mode 100644
index 000000000..76303c7c0
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2697.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2698.png b/iann/train/datasets/val_data/gt/frame_2698.png
new file mode 100644
index 000000000..8d8255b19
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2698.png differ
diff --git a/iann/train/datasets/val_data/gt/frame_2699.png b/iann/train/datasets/val_data/gt/frame_2699.png
new file mode 100644
index 000000000..7e7ccfb30
Binary files /dev/null and b/iann/train/datasets/val_data/gt/frame_2699.png differ
diff --git a/iann/train/datasets/val_data/img/frame_2687.jpg b/iann/train/datasets/val_data/img/frame_2687.jpg
new file mode 100644
index 000000000..924674e6d
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2687.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2689.jpg b/iann/train/datasets/val_data/img/frame_2689.jpg
new file mode 100644
index 000000000..765135b7f
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2689.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2690.jpg b/iann/train/datasets/val_data/img/frame_2690.jpg
new file mode 100644
index 000000000..39274be54
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2690.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2691.jpg b/iann/train/datasets/val_data/img/frame_2691.jpg
new file mode 100644
index 000000000..76f6df3cc
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2691.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2692.jpg b/iann/train/datasets/val_data/img/frame_2692.jpg
new file mode 100644
index 000000000..a823c2e2f
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2692.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2693.jpg b/iann/train/datasets/val_data/img/frame_2693.jpg
new file mode 100644
index 000000000..5ef38e739
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2693.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2695.jpg b/iann/train/datasets/val_data/img/frame_2695.jpg
new file mode 100644
index 000000000..3f664be35
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2695.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2697.jpg b/iann/train/datasets/val_data/img/frame_2697.jpg
new file mode 100644
index 000000000..83e4a9e7c
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2697.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2698.jpg b/iann/train/datasets/val_data/img/frame_2698.jpg
new file mode 100644
index 000000000..8465f582f
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2698.jpg differ
diff --git a/iann/train/datasets/val_data/img/frame_2699.jpg b/iann/train/datasets/val_data/img/frame_2699.jpg
new file mode 100644
index 000000000..99b5e5827
Binary files /dev/null and b/iann/train/datasets/val_data/img/frame_2699.jpg differ
diff --git a/iann/train/ritm_train.py b/iann/train/ritm_train.py
new file mode 100644
index 000000000..1495c0be4
--- /dev/null
+++ b/iann/train/ritm_train.py
@@ -0,0 +1,296 @@
+import os
+import sys
+import random
+import argparse
+from easydict import EasyDict as edict
+
+import paddle
+import paddle.nn as nn
+import paddleseg.transforms as T
+from paddleseg.utils import logger, get_sys_env, logger
+from albumentations import (
+    Compose, ShiftScaleRotate, PadIfNeeded, RandomCrop,
+    RGBShift, RandomBrightnessContrast, RandomRotate90, HorizontalFlip
+)
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
+
+from model.model import get_hrnet_model, DistMapsHRNetModel, get_deeplab_model, get_shufflenet_model
+from model.modeling.hrnet_ocr import HighResolutionNet
+from model.loss import *
+from data.points_sampler import MultiPointSampler
+from data.mdiy import MyDataset
+from util.config import cfgData
+from util.util import *
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config", dest="cfg", type=str, \
+        default='iann/train/train_config.yaml', help='The config file.')
+    return parser.parse_args()
+
+
+def main():
+    env_info = get_sys_env()
+    info = ['{}: {}'.format(k, v) for k, v in env_info.items()]
+    info = '\n'.join(['', format('Environment Information', '-^48s')] + info + ['-' * 48])
+    logger.info(info)
+    place = 'gpu' if env_info['Paddle compiled with cuda'] and \
+            env_info['GPUs used'] else 'cpu'
+    paddle.set_device(place)
+    nranks = paddle.distributed.ParallelEnv().nranks
+    local_rank = paddle.distributed.ParallelEnv().local_rank
+
+    cfg = cfgData(parse_args().cfg)
+    model_cfg = edict()
+    model_cfg.input_normalization = {
+        'mean': [.485, .456, .406],
+        'std': [.229, .224, .225]
+    }
+    model_cfg.num_max_points = 10
+    model_cfg.input_transform = T.Compose([
+        T.Normalize(mean=model_cfg.input_normalization['mean'],
+                    std=model_cfg.input_normalization['std'])],
+                    to_rgb=False)
+    nn.initializer.set_global_initializer(nn.initializer.Normal(), nn.initializer.Constant())
+    models = cfg.get('model')
+    if models.get('type') == 'deeplab':
+        model = get_deeplab_model(backbone=models.get('backbone'), is_ritm=models.get('is_ritm'))
+    elif models.get('type') == 'hrnet':
+        model = get_hrnet_model(width=models.get('width'), ocr_width=models.get('ocr_width'), \
+                                with_aux_output=models.get('with_aux_output'), is_ritm=models.get('is_ritm'))
+    elif models.get('type') == 'shufflenet':
+        model = get_shufflenet_model()
+    if models.get('weights') != 'None':
+        model.load_weights(models.get('weights'))
+    backbone_params, other_params = model.get_trainable_params()
+
+    if nranks > 1:
+        if not paddle.distributed.parallel.parallel_helper._is_parallel_ctx_initialized():
+            paddle.distributed.init_parallel_env()
+            ddp_net = paddle.DataParallel(model)
+        else:
+            ddp_net = paddle.DataParallel(model)
+        train(ddp_net, cfg, model_cfg, backbone_params=backbone_params,
+              other_params=other_params)
+    else:
+        train(model, cfg, model_cfg, backbone_params=backbone_params,
+              other_params=other_params)
+
+
+def train(model, cfg, model_cfg, backbone_params=None, other_params=None):
+    local_rank = paddle.distributed.ParallelEnv().local_rank
+
+    max_iters = cfg.get('iters')
+    save_dir = cfg.get('save_dir')
+    batch_size = cfg.get('batch_size') if cfg.get('batch_size') > 1 else 1
+    val_batch_size = batch_size
+    input_normalization = model_cfg.input_normalization
+    crop_size = cfg.get('train_dataset').get('crop_size')
+    log_iters = cfg.get('log_iters')
+    save_interval = cfg.get('save_interval')
+    num_masks = 1
+
+    train_augmentator = Compose([
+        UniformRandomResize(scale_range=(0.75, 1.40)),
+        HorizontalFlip(),
+        PadIfNeeded(min_height=crop_size[0], min_width=crop_size[1], border_mode=0),
+        RandomCrop(*crop_size),
+        RandomBrightnessContrast(brightness_limit=(-0.25, 0.25), contrast_limit=(-0.15, 0.4), p=0.75),
+        RGBShift(r_shift_limit=10, g_shift_limit=10, b_shift_limit=10, p=0.75)
+    ], p=1.0)
+    val_augmentator = Compose([
+        PadIfNeeded(min_height=crop_size[0], min_width=crop_size[1], border_mode=0),
+        RandomCrop(*crop_size)
+    ], p=1.0)
+
+    def scale_func(image_shape):
+        return random.uniform(0.75, 1.25)
+
+    points_sampler = MultiPointSampler(
+        model_cfg.num_max_points, prob_gamma=0.7,
+        merge_objects_prob=0.15, max_num_merged_objects=2)
+    trainset = MyDataset(
+        dataset_path=cfg.get('dataset').get('dataset_path'),
+        folder_name=cfg.get('train_dataset').get('folder_name'),
+        images_dir_name=cfg.get('dataset').get('image_name'),
+        masks_dir_name=cfg.get('dataset').get('label_name'),
+        num_masks=num_masks,
+        augmentator=train_augmentator,
+        points_from_one_object=False,
+        input_transform=model_cfg.input_transform,
+        min_object_area=80,
+        keep_background_prob=0.0,
+        image_rescale=scale_func,
+        points_sampler=points_sampler,
+        samples_scores_path=None,
+        samples_scores_gamma=1.25
+    )
+    valset = MyDataset(
+        dataset_path=cfg.get('dataset').get('dataset_path'),
+        folder_name=cfg.get('val_dataset').get('folder_name'),
+        images_dir_name=cfg.get('dataset').get('image_name'),
+        masks_dir_name=cfg.get('dataset').get('label_name'),
+        augmentator=val_augmentator,
+        num_masks=num_masks,
+        points_from_one_object=False,
+        input_transform=model_cfg.input_transform,
+        min_object_area=80,
+        image_rescale=scale_func,
+        points_sampler=points_sampler
+    )
+    batch_sampler = paddle.io.DistributedBatchSampler(
+        trainset, batch_size=batch_size, shuffle=True, drop_last=True)
+    loader = paddle.io.DataLoader(
+        trainset,
+        batch_sampler=batch_sampler,
+        return_list=True,
+    )
+    val_batch_sampler = paddle.io.DistributedBatchSampler(
+        valset, batch_size=batch_size, shuffle=True, drop_last=True)
+    val_loader = paddle.io.DataLoader(
+        valset,
+        batch_sampler=val_batch_sampler,
+        return_list=True,
+    )
+
+    if cfg.get('use_vdl'):
+        from visualdl import LogWriter
+        log_writer = LogWriter(save_dir)
+
+    iters_per_epoch = len(batch_sampler)
+
+    opt = None
+    if cfg.get('optimizer').get('type') == 'adam':
+        opt = paddle.optimizer.Adam
+    elif cfg.get('optimizer').get('type') == 'sgd':
+        opt = paddle.optimizer.SGD
+    else:
+        raise ValueError('Opt only have adam or sgd now.')
+    lr = None
+    if cfg.get('learning_rate').get('decay').get('type') == 'poly':
+        lr = paddle.optimizer.lr.PolynomialDecay
+    else:
+        raise ValueError('Lr only have poly now.')
+    optimizer1 = opt(learning_rate=lr(float(cfg.get('learning_rate').get('value_1')),
+        decay_steps=cfg.get('learning_rate').get('decay').get('steps'), 
+        end_lr=cfg.get('learning_rate').get('decay').get('end_lr'), 
+        power=cfg.get('learning_rate').get('decay').get('power')),
+        parameters=other_params)
+    optimizer2 = opt(learning_rate=lr(float(cfg.get('learning_rate').get('value_1')),
+        decay_steps=cfg.get('learning_rate').get('decay').get('steps'), 
+        end_lr=cfg.get('learning_rate').get('decay').get('end_lr'), 
+        power=cfg.get('learning_rate').get('decay').get('power')),
+        parameters=backbone_params)
+    instance_loss = NormalizedFocalLossSigmoid(alpha=0.5, gamma=2)
+    instance_aux_loss = SigmoidBinaryCrossEntropyLoss()
+    model.train()
+    # with open('mobilenet_model.txt', 'w') as f:
+    #     for keys, values in model.state_dict().items():
+    #         f.write(keys +'\t'+str(values.shape)+"\n")
+    iters = 0
+    avg_loss = 0.0
+    while iters < max_iters:
+        for data in loader:
+            iters += 1
+            if iters > max_iters:
+                break
+            if len(data) == 3:
+                images, points, masks = data
+            else:
+                images, points = data
+                masks = None
+            if masks is not None:
+                batch_size, num_points, c, h, w = masks.shape
+                masks = masks.reshape([batch_size * num_points, c, h, w])
+
+            output = batch_forward(model, images, masks, points)
+            # output = model(images, points)
+            # print('instance', output['instances'])
+            # print('mask', masks)
+            
+            loss = instance_loss(output['instances'], masks)
+            if 'instances_aux' in output.keys():
+                aux_loss = instance_aux_loss(output['instances_aux'], masks)
+                total_loss = loss + 0.4 * aux_loss
+            else:
+                total_loss = loss
+            avg_loss += total_loss.numpy()[0]
+            total_loss.backward()
+            optimizer1.step()
+            optimizer2.step()
+            lr = optimizer1.get_lr()
+            if isinstance(optimizer1._learning_rate, paddle.optimizer.lr.LRScheduler):
+                optimizer1._learning_rate.step()
+            if isinstance(optimizer2._learning_rate, paddle.optimizer.lr.LRScheduler):
+                optimizer2._learning_rate.step()
+            model.clear_gradients()
+
+            if iters % log_iters == 0:
+                avg_loss /= log_iters
+                logger.info('Epoch={}, Step={}/{}, loss={:.4f}, lr={}'.format(
+                    (iters - 1) // iters_per_epoch + 1, iters, max_iters, avg_loss, lr))
+                if cfg.get('use_vdl'):
+                    log_writer.add_scalar('Train/loss', avg_loss, iters)
+                    log_writer.add_scalar('Train/lr', lr, iters)
+                avg_loss = 0.0
+            if (iters % save_interval == 0 or iters == max_iters) and local_rank == 0:
+                model.eval()
+                total_len = len(val_loader)
+                val_iou = 0
+                for val_num, val_data in enumerate(val_loader):
+                    if len(data) == 3:
+                        val_images, val_points, val_masks = val_data
+                    else:
+                        val_images, val_points = val_data
+                        val_masks = None
+                    if val_masks is not None:
+                        val_batch_size, val_num_points, val_c, val_h, val_w = val_masks.shape
+                        val_masks = val_masks.reshape([val_batch_size * val_num_points, val_c, val_h, val_w])
+                    val_output = batch_forward(model, val_images, val_masks, val_points, is_train=False)['instances']    
+                    # val_output = model(val_images, val_points)['instances']
+                    # print('max', paddle.max(val_output))
+                    # print('output shape', val_output.shape)
+                    val_output = nn.functional.interpolate(val_output, mode='bilinear', align_corners=True,
+                                               size=val_masks.shape[2:])
+                    val_output = val_output > 0.5
+                    iter_iou = get_iou(val_masks.numpy(), val_output.numpy())
+                    val_iou += iter_iou
+                logger.info('mean iou of iter {} is {}'.format(iters, val_iou / total_len))
+
+                if cfg.get('use_vdl'):
+                    log_writer.add_scalar('Eval/miou', val_iou / total_len, iters)
+
+                current_save_dir = os.path.join(save_dir, "iter_{}".format(iters))
+
+                if not os.path.isdir(current_save_dir):
+                    os.makedirs(current_save_dir)
+                paddle.save(model.state_dict(),
+                            os.path.join(current_save_dir, 'model.pdparams'))
+                model.train()
+
+
+def batch_forward(model, image, gt_mask, points, is_train=True):
+    orig_image, orig_gt_mask = image.clone(), gt_mask.clone()
+    prev_output = paddle.zeros_like(image, dtype='float32')[:, :1, : :]
+    # last_click_indx = None
+    num_iters = random.randint(1, 3)
+    if is_train:
+        model.eval()
+    with paddle.no_grad():
+        for click_indx in range(num_iters):
+            net_input = paddle.concat([image, prev_output], axis=1)
+            prev_output = model(net_input, points)['instances']
+            prev_output = nn.functional.sigmoid(prev_output)
+            points = get_next_points(prev_output, orig_gt_mask, points, click_indx + 1)
+    if is_train:
+        model.train()
+    net_input = paddle.concat([image, prev_output], axis=1)
+    output = model(net_input, points)
+    return output
+
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file
diff --git a/iann/train/train_config.yaml b/iann/train/train_config.yaml
new file mode 100644
index 000000000..f4e7d0b29
--- /dev/null
+++ b/iann/train/train_config.yaml
@@ -0,0 +1,37 @@
+iters: 100000
+batch_size: 16
+save_interval: 1000
+log_iters: 10
+worker: 4
+save_dir: model_output
+use_vdl: False
+
+dataset:
+  dataset_path: iann/train/datasets
+  image_name: img
+  label_name: gt
+
+train_dataset:
+  crop_size: [320, 480]
+  folder_name: train_data
+
+val_dataset:
+  folder_name: val_data
+
+optimizer:
+  type: adam
+
+learning_rate:
+  value_1: 5e-5
+  value_2: 5e-6
+  decay:
+    type: poly
+    steps: 1000
+    power: 0.9
+    end_lr: 0.0
+
+model:
+  type: deeplab
+  backbone: resnet18
+  is_ritm: True
+  weights: None
\ No newline at end of file
diff --git a/iann/ui.py b/iann/ui.py
index c0d851218..b1e4ddce3 100644
--- a/iann/ui.py
+++ b/iann/ui.py
@@ -111,7 +111,7 @@ def setupUi(self, MainWindow):
         self.statusbar.setObjectName("statusbar")
         self.statusbar.setStyleSheet("QStatusBar::item {border: none;}")
         MainWindow.setStatusBar(self.statusbar)
-        self.statusbar.addPermanentWidget(self.show_logo("iann/resources/paddle.png"))
+        self.statusbar.addPermanentWidget(self.show_logo("iann/resource/paddle.png"))
         ## -----
         ## -- 图形区域 --
         ImageRegion = QtWidgets.QHBoxLayout(CentralWidget)
diff --git a/iann/util/config.py b/iann/util/config.py
new file mode 100644
index 000000000..9b48fe5d8
--- /dev/null
+++ b/iann/util/config.py
@@ -0,0 +1,19 @@
+import yaml
+
+
+class cfgData(object):
+    def __init__(self, yaml_file):
+        with open(yaml_file, 'r', encoding="utf-8") as f:
+            fig_data = f.read()
+            self.dicts = yaml.load(fig_data)
+
+    def get(self, key):
+        if key in self.dicts.keys():
+            return self.dicts[key]
+        else:
+            raise ValueError('Not find this keyword.')
+
+
+if __name__ == '__main__':
+    cfg = cfgData('iann/train/train_config.yaml')
+    print(cfg.get('use_vdl'))
\ No newline at end of file