little problem

Author: seanXYZ

dataset.py

@@ -7,6 +7,8 @@
 from PIL import Image
 import torchvision.transforms as transforms
 
+import torchvision.transforms.functional as VF
+
 class SegDataset(Dataset):
     def __init__(self, dataRoot, transforms_= None, transforms_mask = None, subFold="Train_NG", isTrain=True):
 
@@ -32,18 +34,33 @@ def __getitem__(self, index):
 
         idx = index %  self.len
 
-        img  = Image.open(self.imgFiles[idx]).convert("RGB")
-        img = self.transform(img)
+
 
         if self.isTrain==True:
+
+            img  = Image.open(self.imgFiles[idx]).convert("RGB")
+            
             #mask = Image.open(self.labelFiles[idx]).convert("RGB")   
             mat = cv2.imread(self.labelFiles[idx], cv2.IMREAD_GRAYSCALE)
             kernel = np.ones((5, 5), np.uint8)
             matD = cv2.dilate(mat, kernel)
             mask = Image.fromarray(matD)               # image2 is a PIL image    
+
+            if np.random.rand(1) > 0.5:
+                mask = VF.hflip(mask)
+                img  = VF.hflip(img)
+            
+            if np.random.rand(1) > 0.5:
+                mask = VF.vflip(mask)
+                img  = VF.vflip(img)
+
+            img = self.transform(img)
             mask = self.maskTransform(mask)
+
             return {"img":img, "mask":mask}
         else:
+            img  = Image.open(self.imgFiles[idx]).convert("RGB")
+            img = self.transform(img)
             return {"img":img}
 
     def __len__(self):

models.py

@@ -88,29 +88,29 @@ def __init__(self, init_weights=True):
 
         self.layer1 = nn.Sequential(
                             nn.MaxPool2d(2),
-                            nn.Conv2d(1025, 8, 5),
+                            nn.Conv2d(1025, 8, 5, stride=1, padding=2),
                             nn.BatchNorm2d(8),
                             nn.ReLU(inplace=True),
                             nn.MaxPool2d(2),
-                            nn.Conv2d(8, 16, 5),
+                            nn.Conv2d(8, 16, 5, stride=1, padding=2),
                             nn.BatchNorm2d(16),
                             nn.ReLU(inplace=True),
-                            nn.Conv2d(16, 32, 5),
+                            nn.Conv2d(16, 32, 5, stride=1, padding=2),
                             nn.BatchNorm2d(32),
                             nn.ReLU(inplace=True)
                         )
 
         self.fc =  nn.Sequential(
-                            nn.Linear(66, 1, bias=False)
+                            nn.Linear(66, 1, bias=False),
+                            nn.Sigmoid()
                         )
 
         if init_weights == True:
             pass
 
     def forward(self, f, s):
-        x = torch.cat((f, s), 1)
-        x1 = self.layer1(x)
-
+        xx = torch.cat((f, s), 1)
+        x1 = self.layer1(xx)
         x2 = x1.view(x1.size(0), x1.size(1), -1)
         s2 = s.view(s.size(0), s.size(1), -1)
 
@@ -129,12 +129,20 @@ def forward(self, f, s):
 
     snet = SegmentNet()
     dnet = DecisionNet() 
+    img  =  torch.randn(4, 3, 704, 256)
+
+    snet.eval()
 
-    img  =  torch.randn(4, 3, 512, 512)
+    snet = snet.cuda()
+    dnet = dnet.cuda()
+    img = img.cuda()
 
-    f,s = snet.forward(img)
-    c = dnet.forward(f,s)
+    ret = snet(img)
+    f = ret["f"]
+    s = ret["seg"]
 
+    c = dnet(f, s)
+    print(c)
     pass
 
 

test.py

@@ -0,0 +1,92 @@
+from models import SegmentNet, DecisionNet, weights_init_normal
+from dataset import SegDataset
+
+import torch.nn as nn
+import torch
+
+from torchvision import datasets
+from torchvision.utils import save_image
+import torchvision.transforms as transforms
+from torch.autograd import Variable
+from torch.utils.data import DataLoader
+
+import os
+import sys
+import argparse
+import time
+import PIL.Image as Image
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--cuda", type=bool, default=True, help="number of gpu")
+parser.add_argument("--test_seg_epoch", type=int, default=60, help="test segment epoch")
+parser.add_argument("--test_dec_epoch", type=int, default=50, help="test segment epoch")
+parser.add_argument("--img_height", type=int, default=704, help="size of image height")
+parser.add_argument("--img_width", type=int, default=256, help="size of image width")
+
+opt = parser.parse_args()
+
+print(opt)
+
+dataSetRoot = "/home/sean/Data/KolektorSDD_sean"
+
+# ***********************************************************************
+
+# Build nets
+segment_net = SegmentNet(init_weights=True)
+decision_net = DecisionNet(init_weights=True)
+
+if opt.cuda:
+    segment_net = segment_net.cuda()
+    decision_net = decision_net.cuda()
+
+if opt.test_seg_epoch != 0:
+    # Load pretrained models
+    segment_net.load_state_dict(torch.load("./saved_models/segment_net_%d.pth" % (opt.test_seg_epoch)))
+
+if opt.test_dec_epoch != 0:
+    # Load pretrained models
+    decision_net.load_state_dict(torch.load("./saved_models/decision_net_%d.pth" % (opt.test_dec_epoch)))
+
+transforms_ = transforms.Compose([
+    transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),
+    transforms.ToTensor(),
+    #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+])
+
+
+testloader = DataLoader(
+    SegDataset(dataSetRoot, transforms_=transforms_, transforms_mask= None,  subFold="Test", isTrain=False),
+    batch_size=1,
+    shuffle=False,
+    num_workers=0,
+)
+
+segment_net.eval()
+decision_net.eval()
+
+for i, testBatch in enumerate(testloader):
+    t1 = time.time()
+    imgTest = testBatch["img"].cuda()
+    rstTest = segment_net(imgTest)
+
+    fTest = rstTest["f"]
+    segTest = rstTest["seg"]
+
+    cTest = decision_net(fTest, segTest)
+
+    t2 = time.time()
+
+    if cTest.item() > 0.5:
+        labelStr = "NG"
+    else: 
+        labelStr = "OK"
+
+    save_path_str = os.path.join(dataSetRoot, "testResult")
+
+    if os.path.exists(save_path_str) == False:
+        os.makedirs(save_path_str, exist_ok=True)
+
+    print("processing image NO %d, time comsuption %fs"%(i, t2 - t1))
+    save_image(imgTest.data, "%s/img_%d_%s.jpg"% (save_path_str, i, labelStr))
+    save_image(segTest.data, "%s/img_%d_seg_%s.jpg"% (save_path_str, i, labelStr))