--- a
+++ b/CaraNet/train_blood.py
@@ -0,0 +1,251 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Jul 29 17:41:30 2021
+
+@author: angelou
+"""
+
+import torch
+from torch.autograd import Variable
+import os
+import argparse
+from datetime import datetime
+from utils.dataloader import get_loader,test_dataset
+from utils.utils import clip_gradient, adjust_lr, AvgMeter
+import torch.nn.functional as F
+import numpy as np
+from torchstat import stat
+from CaraNet import caranet
+
+
+
+def structure_loss(pred, mask):
+    
+    weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask)
+    wbce = F.binary_cross_entropy_with_logits(pred, mask, reduce='none')
+    wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))
+
+    pred = torch.sigmoid(pred)
+    inter = ((pred * mask)*weit).sum(dim=(2, 3))
+    union = ((pred + mask)*weit).sum(dim=(2, 3))
+    wiou = 1 - (inter + 1)/(union - inter+1)
+    
+    return (wbce + wiou).mean()
+
+
+
+
+
+def test(model, path):
+    
+    ##### put your data_path of TestDataSet/Kvasir here #####
+    data_path = path
+    #########################################################
+    
+    model.eval()
+    image_root = '{}/images/'.format(data_path)
+    gt_root = '{}/masks/'.format(data_path)
+    test_loader = test_dataset(image_root, gt_root, 512)
+    b=0.0
+    print('[test_size]',test_loader.size)
+    for i in range(test_loader.size):
+        image, gt, name = test_loader.load_data()
+        gt = np.asarray(gt, np.float32)
+        gt /= (gt.max() + 1e-8)
+        image = image.cuda()
+        
+        res5,res3,res2,res1 = model(image)
+        res = res5
+        res = F.upsample(res, size=gt.shape, mode='bilinear', align_corners=False)
+        res = res.sigmoid().data.cpu().numpy().squeeze()
+        res = (res - res.min()) / (res.max() - res.min() + 1e-8)
+        
+        input = res
+        target = np.array(gt)
+        N = gt.shape
+        smooth = 1
+        input_flat = np.reshape(input,(-1))
+        target_flat = np.reshape(target,(-1))
+ 
+        intersection = (input_flat*target_flat)
+        
+        loss =  (2 * intersection.sum() + smooth) / (input.sum() + target.sum() + smooth)
+
+        a =  '{:.4f}'.format(loss)
+        a = float(a)
+        b = b + a
+        
+    return b/60
+
+
+
+def train(train_loader, model, optimizer, epoch, test_path):
+    model.train()
+    # ---- multi-scale training ----
+    size_rates = [0.75, 1, 1.25]
+    loss_record1, loss_record2, loss_record3, loss_record5 = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
+    for i, pack in enumerate(train_loader, start=1):
+        for rate in size_rates:
+            optimizer.zero_grad()
+            # ---- data prepare ----
+            images, gts = pack
+            images = Variable(images).cuda()
+            gts = Variable(gts).cuda()
+            # ---- rescale ----
+            trainsize = int(round(opt.trainsize*rate/32)*32)
+            if rate != 1:
+                images = F.upsample(images, size=(trainsize, trainsize), mode='bilinear', align_corners=True)
+                gts = F.upsample(gts, size=(trainsize, trainsize), mode='bilinear', align_corners=True)
+            # ---- forward ----
+            lateral_map_5,lateral_map_3,lateral_map_2,lateral_map_1 = model(images)
+            # ---- loss function ----
+            loss5 = structure_loss(lateral_map_5, gts)
+            loss3 = structure_loss(lateral_map_3, gts)
+            loss2 = structure_loss(lateral_map_2, gts)
+            loss1 = structure_loss(lateral_map_1, gts)
+            
+            
+            loss = loss5 +loss3 + loss2 + loss1
+            # ---- backward ----
+            loss.backward()
+            clip_gradient(optimizer, opt.clip)
+            optimizer.step()
+            # ---- recording loss ----
+            if rate == 1:
+                
+                loss_record5.update(loss5.data, opt.batchsize)
+                loss_record3.update(loss3.data, opt.batchsize)
+                loss_record2.update(loss2.data, opt.batchsize)
+                loss_record1.update(loss1.data, opt.batchsize)
+        # ---- train visualization ----
+        if i % 20 == 0 or i == total_step:
+            print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], '
+                  ' lateral-5: {:0.4f}], lateral-3: {:0.4f}], lateral-2: {:0.4f}], lateral-1: {:0.4f}]'.
+                  format(datetime.now(), epoch, opt.epoch, i, total_step,
+                          loss_record5.show(),loss_record3.show(),loss_record2.show(),loss_record1.show()))
+    save_path = 'snapshots/{}/'.format(opt.train_save)
+    os.makedirs(save_path, exist_ok=True)
+    
+    
+    
+    
+    
+    if (epoch+1) % 1 == 0:
+        meandice = test(model,test_path)
+        
+        fp = open('log/log.txt','a')
+        fp.write(str(meandice)+'\n')
+        fp.close()
+        
+        fp = open('log/best.txt','r')
+        best = fp.read()
+        fp.close()
+        
+        if meandice > float(best):
+            fp = open('log/best.txt','w')
+            fp.write(str(meandice))
+            fp.close()
+            # best = meandice
+            fp = open('log/best.txt','r')
+            best = fp.read()
+            fp.close()
+            torch.save(model.state_dict(), save_path + 'CaraNet-best.pth' )
+            print('[Saving Snapshot:]', save_path + 'CaraNet-best.pth',meandice,'[best:]',best)
+            
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument('--epoch', type=int,
+                        default=10, help='epoch number')
+    
+    parser.add_argument('--lr', type=float,
+                        default=1e-4, help='learning rate')
+    
+    parser.add_argument('--optimizer', type=str,
+                        default='Adam', help='choosing optimizer Adam or SGD')
+    
+    parser.add_argument('--augmentation',
+                        default=False, help='choose to do random flip rotation')
+    
+    parser.add_argument('--batchsize', type=int,
+                        default=6, help='training batch size')
+    
+    parser.add_argument('--trainsize', type=int,
+                        default=352, help='training dataset size')
+    
+    parser.add_argument('--clip', type=float,
+                        default=0.5, help='gradient clipping margin')
+    
+    parser.add_argument('--decay_rate', type=float,
+                        default=0.1, help='decay rate of learning rate')
+    
+    parser.add_argument('--decay_epoch', type=int,
+                        default=50, help='every n epochs decay learning rate')
+    
+    parser.add_argument('--train_path', type=str,
+                        default='/home/data/spleen_blood/CaraNet/TrainDataset/train/', help='path to train dataset')
+    
+    parser.add_argument('--test_path', type=str,
+                        default='/home/data/spleen_blood/CaraNet/TestDataset/test/' , help='path to testing Kvasir dataset')
+    
+    parser.add_argument('--train_save', type=str,
+                        default='')
+    
+    opt = parser.parse_args()
+
+    # ---- build models ----
+    torch.cuda.set_device(4)  # set your gpu device
+    model = caranet().cuda()
+    # ---- flops and params ----
+
+    # from utils.utils import CalParams
+    # x = torch.randn(1, 3, 352, 352).cuda()
+    # CalParams(model, x)
+
+    params = model.parameters()
+    
+    if opt.optimizer == 'Adam':
+        optimizer = torch.optim.Adam(params, opt.lr)
+    else:
+        optimizer = torch.optim.SGD(params, opt.lr, weight_decay = 1e-4, momentum = 0.9)
+        
+    print(optimizer)
+
+    image_root = '{}/image/'.format(opt.train_path)
+    gt_root = '{}/mask/'.format(opt.train_path)
+
+    train_loader = get_loader(image_root, gt_root, batchsize=opt.batchsize, trainsize=opt.trainsize, augmentation = opt.augmentation)
+    total_step = len(train_loader)
+
+    print("#"*20, "Start Training", "#"*20)
+
+    for epoch in range(1, opt.epoch):
+        adjust_lr(optimizer, opt.lr, epoch, 0.1, 200)
+        train(train_loader, model, optimizer, epoch, opt.test_path)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+