normalize input

Author: huochaitiantang

core/data.py

@@ -71,9 +71,10 @@ def get_files(mydir):
 
 # Dataset not composite online
 class MatDatasetOffline(torch.utils.data.Dataset):
-    def __init__(self, args, transform=None):
+    def __init__(self, args, transform=None, normalize=None):
         self.samples=[]
         self.transform = transform
+        self.normalize = normalize
         self.args = args
         self.size_h = args.size_h
         self.size_w = args.size_w
@@ -144,6 +145,14 @@ def __getitem__(self,index):
         trimap = gen_trimap(alpha)
         grad = compute_gradient(img)
 
+        if self.normalize:
+            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            # first, 0-255 to 0-1
+            # second, x-mean/std and HWC to CHW
+            img_norm = self.normalize(img_rgb)
+        else:
+            img_norm = None
+
         #img_id = img_info[0].split('/')[-1]
         #cv2.imwrite("result/debug/{}_img.png".format(img_id), img)
         #cv2.imwrite("result/debug/{}_alpha.png".format(img_id), alpha)
@@ -158,7 +167,7 @@ def __getitem__(self,index):
         fg = torch.from_numpy(fg.astype(np.float32)).permute(2, 0, 1)
         bg = torch.from_numpy(bg.astype(np.float32)).permute(2, 0, 1)
 
-        return img, alpha, fg, bg, trimap, grad, img_info
+        return img, alpha, fg, bg, trimap, grad, img_norm, img_info
 
     def __len__(self):
         return len(self.samples)

core/deploy.py

@@ -62,8 +62,18 @@ def inference_once(args, model, scale_img, scale_trimap, aligned=True):
         assert(scale_img.shape[0] == args.size_h)
         assert(scale_img.shape[1] == args.size_w)
 
+    normalize = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean = [0.485, 0.456, 0.406],std = [0.229, 0.224, 0.225])
+    ])
+
+    scale_img_rgb = cv2.cvtColor(scale_img, cv2.COLOR_BGR2RGB)
+    # first, 0-255 to 0-1
+    # second, x-mean/std and HWC to CHW
+    tensor_img = normalize(scale_img_rgb).unsqueeze(0)
+
     scale_grad = compute_gradient(scale_img)
-    tensor_img = torch.from_numpy(scale_img.astype(np.float32)[np.newaxis, :, :, :]).permute(0, 3, 1, 2)
+    #tensor_img = torch.from_numpy(scale_img.astype(np.float32)[np.newaxis, :, :, :]).permute(0, 3, 1, 2)
     tensor_trimap = torch.from_numpy(scale_trimap.astype(np.float32)[np.newaxis, np.newaxis, :, :])
     tensor_grad = torch.from_numpy(scale_grad.astype(np.float32)[np.newaxis, np.newaxis, :, :])
 
@@ -73,7 +83,7 @@ def inference_once(args, model, scale_img, scale_trimap, aligned=True):
         tensor_grad = tensor_grad.cuda()
     #print('Img Shape:{} Trimap Shape:{}'.format(img.shape, trimap.shape))
 
-    input_t = torch.cat((tensor_img, tensor_trimap), 1)
+    input_t = torch.cat((tensor_img, tensor_trimap / 255.), 1)
 
     # forward
     if args.stage <= 1:

core/train.py

@@ -59,7 +59,12 @@ def get_dataset(args):
     args.crop_h = [int(i) for i in args.crop_h.split(',')]
     args.crop_w = [int(i) for i in args.crop_w.split(',')]
 
-    train_set = MatDatasetOffline(args, train_transform)
+    normalize = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean = [0.485, 0.456, 0.406],std = [0.229, 0.224, 0.225])
+    ])
+
+    train_set = MatDatasetOffline(args, train_transform, normalize)
     train_loader = DataLoader(dataset=train_set, num_workers=args.threads, batch_size=args.batchSize, shuffle=True)
 
     return train_loader
@@ -173,6 +178,7 @@ def train(args, model, optimizer, train_loader, epoch):
         fg = Variable(batch[2])
         bg = Variable(batch[3])
         trimap = Variable(batch[4])
+        img_norm = Variable(batch[6])
         img_info = batch[-1]
 
         if args.cuda:
@@ -181,14 +187,15 @@ def train(args, model, optimizer, train_loader, epoch):
             fg = fg.cuda()
             bg = bg.cuda()
             trimap = trimap.cuda()
+            img_norm = img_norm.cuda()
 
         #print("Shape: Img:{} Alpha:{} Fg:{} Bg:{} Trimap:{}".format(img.shape, alpha.shape, fg.shape, bg.shape, trimap.shape))
         #print("Val: Img:{} Alpha:{} Fg:{} Bg:{} Trimap:{} Img_info".format(img, alpha, fg, bg, trimap, img_info))
 
         adjust_learning_rate(args, optimizer, epoch)
         optimizer.zero_grad()
 
-        pred_mattes, pred_alpha = model(torch.cat((img, trimap), 1))
+        pred_mattes, pred_alpha = model(torch.cat((img_norm, trimap / 255.), 1))
 
         if args.stage == 0:
             # stage0 loss, simple alpha loss

deploy.sh

@@ -1,6 +1,6 @@
 #/bin/bash
 
-TEST_DATA_ROOT=/data/datasets/matting/Combined_Dataset/Test_set/comp
+TEST_DATA_ROOT=/home/liuliang/DISK_2T/datasets/matting/Combined_Dataset/Test_set/comp
 
 python core/deploy.py \
     --size_h=320 \
@@ -9,7 +9,7 @@ python core/deploy.py \
     --trimapDir=$TEST_DATA_ROOT/trimap \
     --alphaDir=$TEST_DATA_ROOT/alpha \
     --saveDir=result/stage0 \
-    --resume=model/stage0/ckpt_e19.pth \
+    --resume=model/stage0_norm/ckpt_e1.pth \
     --cuda \
     --stage=0 \
     --crop_or_resize=whole \

train.sh

@@ -1,5 +1,5 @@
 #/bin/bash
-DATA_ROOT=/data/datasets/matting/Combined_Dataset
+DATA_ROOT=/home/liuliang/DISK_2T/datasets/matting/Combined_Dataset
 TRAIN_DATA_ROOT=$DATA_ROOT/Training_set/comp
 TEST_DATA_ROOT=$DATA_ROOT/Test_set/comp
 
@@ -12,14 +12,14 @@ python core/train.py \
     --fgDir=$TRAIN_DATA_ROOT/fg \
     --bgDir=$TRAIN_DATA_ROOT/bg \
     --imgDir=$TRAIN_DATA_ROOT/image \
-    --saveDir=model/stage0 \
+    --saveDir=model/stage0_norm \
     --batchSize=1 \
     --nEpochs=25 \
     --step=-1 \
     --lr=0.00001 \
     --wl_weight=0.5 \
     --threads=4 \
-    --printFreq=1 \
+    --printFreq=10 \
     --ckptSaveFreq=1 \
     --cuda \
     --stage=0 \
@@ -30,5 +30,5 @@ python core/train.py \
     --testAlphaDir=$TEST_DATA_ROOT/alpha \
     --testResDir=result/tmp \
     --crop_or_resize=whole \
-    --max_size=1600
-    #--resume=model/stage0/ckpt_e6.pth \
+    --max_size=1600 \
+    #--resume=model/stage0_norm/ckpt_e2.pth \