update STDC

conf/model.json

@@ -1,6 +1,6 @@
 {
 
-	"gpu_process":{"det_weights":"weights/yolov5/class5/best_5classes.pt","seg_nclass":2,"seg_weights": "weights/BiSeNet/checkpoint.pth"    },
+	"gpu_process":{"det_weights":"weights/yolov5/class5/best_5classes.pt","seg_nclass":2,"seg_weights_1": "weights/BiSeNet/checkpoint.pth","seg_weights":"weights/STDC/model_maxmIOU75_1720_0.946_360640.pth"    },
 
 	"post_process":{ "name":"post_process","conf_thres":0.25,"iou_thres":0.45,"classes":5,"labelnames":"weights/yolov5/class5/labelnames.json","fpsample":240,"debug":false , "rainbows":[ [0,0,255],[0,255,0],[255,0,0],[255,0,255],[255,255,0],[255,129,0],[255,0,127],[127,255,0],[0,255,127],[0,127,255],[127,0,255],[255,127,255],[255,255,127],[127,255,255],[0,255,255],[255,127,255],[127,255,255],  [0,127,0],[0,0,127],[0,255,255]],"outImaDir":"problems/images_tmp","outVideoDir":"problems/videos_save"     },
 

producer.py

@@ -7,7 +7,7 @@ import random,string
 def producer_demo():
 
 
-    cnt_online=1;cnt_offline=0
+    cnt_online=0;cnt_offline=1
     Tecent=False;
     
     #topic_on='thsw';topic_off='thsw2';

segutils/model_stages.py

@@ -0,0 +1,337 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import time
+from stdcnet import STDCNet1446, STDCNet813
+#from models_725.bn import InPlaceABNSync as BatchNorm2d
+# BatchNorm2d = nn.BatchNorm2d
+
+class ConvBNReLU(nn.Module):
+    def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
+        super(ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(in_chan,
+                out_chan,
+                kernel_size = ks,
+                stride = stride,
+                padding = padding,
+                bias = False)
+        # self.bn = BatchNorm2d(out_chan)
+        # self.bn = BatchNorm2d(out_chan, activation='none')
+        self.bn = nn.BatchNorm2d(out_chan)
+        self.relu = nn.ReLU()
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+
+class BiSeNetOutput(nn.Module):
+    def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):
+        super(BiSeNetOutput, self).__init__()
+        self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
+        self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.conv_out(x)
+        return x
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):######################1
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class AttentionRefinementModule(nn.Module):
+    def __init__(self, in_chan, out_chan, *args, **kwargs):
+        super(AttentionRefinementModule, self).__init__()
+        self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
+        self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)
+        # self.bn_atten = nn.BatchNorm2d(out_chan)
+        # self.bn_atten = BatchNorm2d(out_chan, activation='none')
+        self.bn_atten = nn.BatchNorm2d(out_chan)########################2
+
+        self.sigmoid_atten = nn.Sigmoid()
+        self.init_weight()
+
+    def forward(self, x):
+        feat = self.conv(x)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv_atten(atten)
+        atten = self.bn_atten(atten)
+        atten = self.sigmoid_atten(atten)
+        out = torch.mul(feat, atten)
+        return out
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+
+class ContextPath(nn.Module):
+    def __init__(self, backbone='CatNetSmall', pretrain_model='', use_conv_last=False, *args, **kwargs):
+        super(ContextPath, self).__init__()
+        
+        self.backbone_name = backbone
+        if backbone == 'STDCNet1446':
+            self.backbone = STDCNet1446(pretrain_model=pretrain_model, use_conv_last=use_conv_last)
+            self.arm16 = AttentionRefinementModule(512, 128)
+            inplanes = 1024
+            if use_conv_last:
+                inplanes = 1024
+            self.arm32 = AttentionRefinementModule(inplanes, 128)
+            self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+            self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+            self.conv_avg = ConvBNReLU(inplanes, 128, ks=1, stride=1, padding=0)
+
+        elif backbone == 'STDCNet813':
+            self.backbone = STDCNet813(pretrain_model=pretrain_model, use_conv_last=use_conv_last)
+            self.arm16 = AttentionRefinementModule(512, 128)
+            inplanes = 1024
+            if use_conv_last:
+                inplanes = 1024
+            self.arm32 = AttentionRefinementModule(inplanes, 128)
+            self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+            self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+            self.conv_avg = ConvBNReLU(inplanes, 128, ks=1, stride=1, padding=0)
+        else:
+            print("backbone is not in backbone lists")
+            exit(0)
+
+        self.init_weight()
+
+    def forward(self, x):
+        H0, W0 = x.size()[2:]
+
+        feat2, feat4, feat8, feat16, feat32 = self.backbone(x)
+        H8, W8 = feat8.size()[2:]
+        H16, W16 = feat16.size()[2:]
+        H32, W32 = feat32.size()[2:]
+        
+        avg = F.avg_pool2d(feat32, feat32.size()[2:])
+
+        avg = self.conv_avg(avg)
+        avg_up = F.interpolate(avg, (H32, W32), mode='nearest')
+
+        feat32_arm = self.arm32(feat32)
+        feat32_sum = feat32_arm + avg_up
+        feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
+        feat32_up = self.conv_head32(feat32_up)
+
+        feat16_arm = self.arm16(feat16)
+        feat16_sum = feat16_arm + feat32_up
+        feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
+        feat16_up = self.conv_head16(feat16_up)
+        
+        return feat2, feat4, feat8, feat16, feat16_up, feat32_up # x8, x16
+
+        # return feat8, feat16_up  # x8, x16
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):#################3
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class FeatureFusionModule(nn.Module):
+    def __init__(self, in_chan, out_chan, *args, **kwargs):
+        super(FeatureFusionModule, self).__init__()
+        self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
+        self.conv1 = nn.Conv2d(out_chan,
+                out_chan//4,
+                kernel_size = 1,
+                stride = 1,
+                padding = 0,
+                bias = False)
+        self.conv2 = nn.Conv2d(out_chan//4,
+                out_chan,
+                kernel_size = 1,
+                stride = 1,
+                padding = 0,
+                bias = False)
+        self.relu = nn.ReLU(inplace=True)
+        self.sigmoid = nn.Sigmoid()
+        self.init_weight()
+
+    def forward(self, fsp, fcp):
+        fcat = torch.cat([fsp, fcp], dim=1)
+        feat = self.convblk(fcat)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv1(atten)
+        atten = self.relu(atten)
+        atten = self.conv2(atten)
+        atten = self.sigmoid(atten)
+        feat_atten = torch.mul(feat, atten)
+        feat_out = feat_atten + feat
+        return feat_out
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):##################4
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class BiSeNet(nn.Module):
+    def __init__(self, backbone, n_classes, pretrain_model='', use_boundary_2=False, use_boundary_4=False,
+                 use_boundary_8=False, use_boundary_16=False, use_conv_last=False):
+        super(BiSeNet, self).__init__()
+        
+        self.use_boundary_2 = use_boundary_2
+        self.use_boundary_4 = use_boundary_4
+        self.use_boundary_8 = use_boundary_8
+        self.use_boundary_16 = use_boundary_16
+        # self.heat_map = heat_map
+        self.cp = ContextPath(backbone, pretrain_model, use_conv_last=use_conv_last)
+        
+        if backbone == 'STDCNet1446':
+            conv_out_inplanes = 128
+            sp2_inplanes = 32
+            sp4_inplanes = 64
+            sp8_inplanes = 256
+            sp16_inplanes = 512
+            inplane = sp8_inplanes + conv_out_inplanes
+
+        elif backbone == 'STDCNet813':
+            conv_out_inplanes = 128
+            sp2_inplanes = 32
+            sp4_inplanes = 64
+            sp8_inplanes = 256
+            sp16_inplanes = 512
+            inplane = sp8_inplanes + conv_out_inplanes
+
+        else:
+            print("backbone is not in backbone lists")
+            exit(0)
+
+        self.ffm = FeatureFusionModule(inplane, 256)
+        self.conv_out = BiSeNetOutput(256, 256, n_classes)
+        self.conv_out16 = BiSeNetOutput(conv_out_inplanes, 64, n_classes)
+        self.conv_out32 = BiSeNetOutput(conv_out_inplanes, 64, n_classes)
+
+        self.conv_out_sp16 = BiSeNetOutput(sp16_inplanes, 64, 1)
+        
+        self.conv_out_sp8 = BiSeNetOutput(sp8_inplanes, 64, 1)
+        self.conv_out_sp4 = BiSeNetOutput(sp4_inplanes, 64, 1)
+        self.conv_out_sp2 = BiSeNetOutput(sp2_inplanes, 64, 1)
+        self.init_weight()
+
+    def forward(self, x):
+        H, W = x.size()[2:]
+        # time_0 = time.time()
+        # feat_res2, feat_res4, feat_res8, feat_res16, feat_cp8, feat_cp16 = self.cp(x)
+        feat_res2, feat_res4, feat_res8, feat_res16, feat_cp8, feat_cp16 = self.cp(x)
+        # print('----backbone', (time.time() - time_0) * 1000)
+        # feat_out_sp2 = self.conv_out_sp2(feat_res2)
+        #
+        # feat_out_sp4 = self.conv_out_sp4(feat_res4)
+        #
+        # feat_out_sp8 = self.conv_out_sp8(feat_res8)
+        #
+        # feat_out_sp16 = self.conv_out_sp16(feat_res16)
+        # time_1 = time.time()
+        feat_fuse = self.ffm(feat_res8, feat_cp8)
+        # print('----ffm', (time.time() - time_1) * 1000)
+        # time_2 = time.time()
+        feat_out = self.conv_out(feat_fuse)
+        # feat_out16 = self.conv_out16(feat_cp8)
+        # feat_out32 = self.conv_out32(feat_cp16)
+
+        feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
+        # print('----conv_out', (time.time() - time_2) * 1000)
+        # feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
+        # feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
+
+
+        # if self.use_boundary_2 and self.use_boundary_4 and self.use_boundary_8:
+        #     return feat_out, feat_out16, feat_out32, feat_out_sp2, feat_out_sp4, feat_out_sp8
+        #
+        # if (not self.use_boundary_2) and self.use_boundary_4 and self.use_boundary_8:
+        #     return feat_out, feat_out16, feat_out32, feat_out_sp4, feat_out_sp8
+        #
+        # if (not self.use_boundary_2) and (not self.use_boundary_4) and self.use_boundary_8:
+        return feat_out
+        
+        # if (not self.use_boundary_2) and (not self.use_boundary_4) and (not self.use_boundary_8):
+        #     return feat_out, feat_out16, feat_out32
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
+        for name, child in self.named_children():
+            child_wd_params, child_nowd_params = child.get_params()
+            if isinstance(child, (FeatureFusionModule, BiSeNetOutput)):
+                lr_mul_wd_params += child_wd_params
+                lr_mul_nowd_params += child_nowd_params
+            else:
+                wd_params += child_wd_params
+                nowd_params += child_nowd_params
+        return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params
+
+
+if __name__ == "__main__":
+    
+    net = BiSeNet('STDCNet813', 19)
+    net.cuda()
+    net.eval()
+    in_ten = torch.randn(1, 3, 768, 1536).cuda()
+    out, out16, out32 = net(in_ten)
+    print(out.shape)
+    # torch.save(net.state_dict(), 'STDCNet813.pth')###
+
+    

segutils/segWaterBuilding.py

@@ -6,7 +6,7 @@ import numpy as np
 from core.models.dinknet import DinkNet34     
 import matplotlib.pyplot as plt           
 
-import matplotlib.pyplot as plt           
+         
 import time
 class SegModel(object):
     def __init__(self, nclass=2,model = None,weights=None,modelsize=512,device='cuda:3',multiOutput=False):
@@ -328,19 +328,20 @@ def illBuildings(pred,image_array0):
     return image_array0,water
              
 def test_water_building_seperately():
-    from core.models.dinknet import DinkNet34_MultiOutput 
+    #from core.models.dinknet import DinkNet34_MultiOutput 
     #create_model('DinkNet34_MultiOutput',[2,5])
     
     image_url = 'temp_pics/DJI_0645.JPG'
     nclass = [2,2]
     outresult=True
-    weights = 'runs/thWaterBuilding_seperate/BiSeNet_MultiOutput/train/experiment_0/checkpoint.pth'
+    weights = '../weights/BiSeNet/checkpoint.pth'
     model =  BiSeNet_MultiOutput(nclass)
     outdir='temp'
     image_dir = '/home/thsw/WJ/data/river_buildings/'
     #image_dir = '/home/thsw/WJ/data/THWaterBuilding/val/images'
     image_url_list=glob.glob('%s/*'%(image_dir))
-    segmodel = SegModel(model=model,nclass=nclass,weights=weights,device='cuda:1',multiOutput=True)
+    #segmodel = SegModel(model=model,nclass=nclass,weights=weights,device='cuda:1',multiOutput=True)
+    segmodel = SegModel(nclass=nclass,weights=weights,device='cuda:1')
     
     print('###line307 image cnt:',len(image_url_list))
     for i,image_url in enumerate(image_url_list[0:1]) :

segutils/segmodel.py

@@ -1,140 +1,128 @@
-import torch
-import sys,os
-sys.path.extend(['segutils'])
-from core.models.bisenet import BiSeNet
-from torchvision import transforms
-import cv2,glob
-import numpy as np
-from core.models.dinknet import DinkNet34     
-import matplotlib.pyplot as plt           
-import time
-class SegModel(object):
-    def __init__(self, nclass=2,weights=None,modelsize=512,device='cuda:0'):
-        #self.args = args       
-        self.model = BiSeNet(nclass)
-        #self.model = DinkNet34(nclass)
-        checkpoint = torch.load(weights)
-        self.modelsize = modelsize
-        self.model.load_state_dict(checkpoint['model'])
-        self.device = device
-        self.model= self.model.to(self.device)
-        '''self.composed_transforms = transforms.Compose([
-            
-            transforms.Normalize(mean=(0.335, 0.358, 0.332), std=(0.141, 0.138, 0.143)),
-            transforms.ToTensor()])   '''
-        self.mean = (0.335, 0.358, 0.332)
-        self.std =   (0.141, 0.138, 0.143)  
-    def eval(self,image):	
-        time0 = time.time()
-        imageH,imageW,imageC = image.shape        
-        image = self.preprocess_image(image)
-        time1 = time.time()
-        self.model.eval()
-        image  = image.to(self.device)
-        with torch.no_grad():
-            output = self.model(image)
-
-        time2 = time.time()    
-        pred = output.data.cpu().numpy()
-        pred = np.argmax(pred, axis=1)[0]#得到每行
-        time3 = time.time()        
-        pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH))
-        time4 = time.time()
-        outstr= 'pre-precess:%.1f ,infer:%.1f ,post-precess:%.1f ,post-resize:%.1f, total:%.1f \n '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
-
-        #print('pre-precess:%.1f ,infer:%.1f ,post-precess:%.1f ,post-resize:%.1f, total:%.1f '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) ))
-        return pred,outstr 
-    def get_ms(self,t1,t0):
-        return (t1-t0)*1000.0            
-    def preprocess_image(self,image):
-        
-        time0 = time.time()
-        image = cv2.resize(image,(self.modelsize,self.modelsize))
-        time0 = time.time()
-        image = image.astype(np.float32)        
-        image /= 255.0
-
-        image[:,:,0] -=self.mean[0]
-        image[:,:,1] -=self.mean[1]
-        image[:,:,2] -=self.mean[2]
-
-        image[:,:,0] /= self.std[0]
-        image[:,:,1] /= self.std[1]
-        image[:,:,2] /= self.std[2]
-        image = cv2.cvtColor( image,cv2.COLOR_RGB2BGR)
-        #image -= self.mean
-        #image /= self.std
-        image = np.transpose(image, ( 2, 0, 1))
-        
-        image = torch.from_numpy(image).float()
-        image = image.unsqueeze(0)
-        
-        
-        return image
-        
-def get_ms(t1,t0):
-    return (t1-t0)*1000.0                  
-    
-     
-def get_largest_contours(contours):
-    areas = [cv2.contourArea(x) for x in contours]
-    max_area = max(areas)
-    max_id = areas.index(max_area)
-
-    return max_id
-    
-if __name__=='__main__':
-    image_url = '/home/thsw2/WJ/data/THexit/val/images/DJI_0645.JPG'
-    nclass = 2
-    weights = '../weights/segmentation/BiSeNet/checkpoint.pth'
-
-    segmodel = SegModel(nclass=nclass,weights=weights)
-
-    image_urls=glob.glob('/home/thsw2/WJ/data/THexit/val/images/*')
-    out_dir ='../runs/detect/exp2-seg';os.makedirs(out_dir,exist_ok=True)
-    for image_url in image_urls[0:1]:
-        image_url = '/home/thsw2/WJ/data/THexit/val/images/54(199).JPG'
-        image_array0 = cv2.imread(image_url)       
-        pred = segmodel.eval(image_array0 )
-
-        #plt.figure(1);plt.imshow(pred);
-        #plt.show()
-        binary0 = pred.copy()
-
-        
-        time0 = time.time()
-        contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
-        max_id = -1
-        if len(contours)>0:
-            max_id = get_largest_contours(contours)
-            binary0[:,:] = 0
-            print(contours[0].shape,contours[1].shape,contours[0])
-            cv2.fillPoly(binary0, [contours[max_id][:,0,:]], 1) 
-      
-        time1 = time.time()
-        
-        #num_labels,_,Areastats,centroids = cv2.connectedComponentsWithStats(binary0,connectivity=4) 
-        time2 = time.time()
-        
-        cv2.drawContours(image_array0,contours,max_id,(0,255,255),3)
-        time3 = time.time()
-        out_url='%s/%s'%(out_dir,os.path.basename(image_url))
-        ret = cv2.imwrite(out_url,image_array0)
-        time4 = time.time()
-        
-        print('image:%s findcontours:%.1f ms , connect:%.1f ms ,draw:%.1f save:%.1f'%(os.path.basename(image_url),get_ms(time1,time0),get_ms(time2,time1), get_ms(time3,time2),get_ms(time4,time3), ) )
-        plt.figure(0);plt.imshow(pred)
-        plt.figure(1);plt.imshow(image_array0)
-        plt.figure(2);plt.imshow(binary0)
-        plt.show()
-        
-        #print(out_url,ret)
-
-
-    
-    
-
-
-
-   
-
+import torch
+from torchvision import transforms
+import cv2,os,sys
+sys.path.extend(['segutils'])
+import numpy as np
+import matplotlib.pyplot as plt
+from model_stages import BiSeNet
+import torch.nn.functional as F
+import time
+
+
+class SegModel(object):
+    def __init__(self, nclass=2,weights=None,modelsize=512,device='cuda:0'):             
+
+        self.model = BiSeNet(backbone='STDCNet813', n_classes=nclass,
+                                use_boundary_2=False, use_boundary_4=False,
+                                use_boundary_8=True, use_boundary_16=False,
+                                use_conv_last=False)
+        self.device = device      
+        self.model.load_state_dict(torch.load(weights, map_location=torch.device(self.device)  ))
+        self.model= self.model.to(self.device)
+        self.mean = (0.485, 0.456, 0.406)
+        self.std = (0.229, 0.224, 0.225)
+
+    def eval(self, image=None):
+        time0 = time.time()
+        imageH, imageW, _ = image.shape
+        image = self.RB_convert(image)
+        img = self.preprocess_image(image)
+        if self.device != 'cpu':
+            imgs = img.to(self.device)
+        else:imgs=img    
+        time1 = time.time()
+        self.model.eval()
+        with torch.no_grad():
+            output = self.model(imgs)
+            
+        time2 = time.time()  
+        pred = output.data.cpu().numpy()
+        pred = np.argmax(pred, axis=1)[0]#得到每行
+        time3 = time.time() 
+        pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH))
+        time4 = time.time()
+        outstr= 'pre-precess:%.1f ,infer:%.1f ,post-cpu-argmax:%.1f ,post-resize:%.1f, total:%.1f \n '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
+       
+
+        return pred, outstr
+
+    def preprocess_image(self, image):
+        image = cv2.resize(image, (640,360), interpolation=cv2.INTER_LINEAR)
+        image = image.astype(np.float32)
+        image /= 255.0
+
+        image[:, :, 0] -= self.mean[0]
+        image[:, :, 1] -= self.mean[1]
+        image[:, :, 2] -= self.mean[2]
+
+        image[:, :, 0] /= self.std[0]
+        image[:, :, 1] /= self.std[1]
+        image[:, :, 2] /= self.std[2]
+
+        image = np.transpose(image, (2, 0, 1))
+        image = torch.from_numpy(image).float()
+        image = image.unsqueeze(0)
+
+        return image
+    def get_ms(self,t1,t0):
+        return (t1-t0)*1000.0     
+    def RB_convert(self,image):
+        image_c = image.copy()
+        image_c[:,:,0] = image[:,:,2]
+        image_c[:,:,2] = image[:,:,0]
+        return image_c
+
+def get_ms(t1,t0):
+    return (t1-t0)*1000.0
+def get_largest_contours(contours):
+    areas = [cv2.contourArea(x) for x in contours]
+    max_area = max(areas)
+    max_id = areas.index(max_area)
+    return max_id
+
+if __name__=='__main__':
+    impth = '/home/thsw/WJ/src/STDC/images/examples'
+    outpth= '/home/thsw/WJ/src/STDC/images/results'
+    folders = os.listdir(impth)
+    weights = '/home/thsw/WJ/src/STDC/model_maxmIOU75_1720_0.946_360640.pth'
+    segmodel = SegModel(nclass=2,weights=weights)
+    
+    for i in range(len(folders)):  
+       
+        imgpath = os.path.join(impth, folders[i])
+        time0 = time.time()
+        
+        #img = Image.open(imgpath).convert('RGB')
+        img = cv2.imread(imgpath)
+        img = np.array(img)
+        
+        
+        time1 = time.time()
+        pred, outstr = segmodel.eval(image=img)#####
+        time2 = time.time()
+        
+        binary0 = pred.copy()
+        contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
+        time3 = time.time()
+        
+        max_id = -1
+        if len(contours)>0:
+            max_id = get_largest_contours(contours)
+            binary0[:,:] = 0
+            cv2.fillPoly(binary0, [contours[max_id][:,0,:]], 1) 
+            
+        cv2.drawContours(img,contours,max_id,(0,255,255),3)
+        time4 = time.time()
+        
+        
+        #img_n = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
+        cv2.imwrite(  os.path.join( outpth,folders[i] ) ,img )
+        time5 = time.time()
+        print('image:%d ,infer:%.1f ms,findcontours:%.1f ms, draw:%.1f, total:%.1f'%(i,get_ms(time2,time1),get_ms(time3,time2),get_ms(time4,time3),get_ms(time4,time1)))                     
+
+
+    
+    
+        
+        
+ 

segutils/segmodel_BiseNet.py

@@ -0,0 +1,143 @@
+import torch
+import sys,os
+sys.path.extend(['segutils'])
+from core.models.bisenet import BiSeNet
+from torchvision import transforms
+import cv2,glob
+import numpy as np
+from core.models.dinknet import DinkNet34     
+import matplotlib.pyplot as plt           
+import time
+class SegModel(object):
+    def __init__(self, nclass=2,weights=None,modelsize=512,device='cuda:0'):
+        #self.args = args       
+        self.model = BiSeNet(nclass)
+        #self.model = DinkNet34(nclass)
+        checkpoint = torch.load(weights)
+        self.modelsize = modelsize
+        self.model.load_state_dict(checkpoint['model'])
+        self.device = device
+        self.model= self.model.to(self.device)
+        '''self.composed_transforms = transforms.Compose([
+            
+            transforms.Normalize(mean=(0.335, 0.358, 0.332), std=(0.141, 0.138, 0.143)),
+            transforms.ToTensor()])   '''
+        self.mean = (0.335, 0.358, 0.332)
+        self.std =   (0.141, 0.138, 0.143)  
+    def eval(self,image):	
+        time0 = time.time()
+        imageH,imageW,imageC = image.shape        
+        image = self.preprocess_image(image)
+        time1 = time.time()
+        self.model.eval()
+        image  = image.to(self.device)
+        with torch.no_grad():
+            output = self.model(image)
+
+        time2 = time.time()    
+        pred = output.data.cpu().numpy()
+        pred = np.argmax(pred, axis=1)[0]#得到每行
+        time3 = time.time()        
+        pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH))
+        time4 = time.time()
+        outstr= 'pre-precess:%.1f ,infer:%.1f ,post-precess:%.1f ,post-resize:%.1f, total:%.1f \n '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
+
+        #print('pre-precess:%.1f ,infer:%.1f ,post-precess:%.1f ,post-resize:%.1f, total:%.1f '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) ))
+        return pred,outstr 
+    def get_ms(self,t1,t0):
+        return (t1-t0)*1000.0            
+    def preprocess_image(self,image):
+        
+        time0 = time.time()
+        image = cv2.resize(image,(self.modelsize,self.modelsize))
+        time0 = time.time()
+        image = image.astype(np.float32)        
+        image /= 255.0
+
+        image[:,:,0] -=self.mean[0]
+        image[:,:,1] -=self.mean[1]
+        image[:,:,2] -=self.mean[2]
+
+        image[:,:,0] /= self.std[0]
+        image[:,:,1] /= self.std[1]
+        image[:,:,2] /= self.std[2]
+        image = cv2.cvtColor( image,cv2.COLOR_RGB2BGR)
+        #image -= self.mean
+        #image /= self.std
+        image = np.transpose(image, ( 2, 0, 1))
+        
+        image = torch.from_numpy(image).float()
+        image = image.unsqueeze(0)
+        
+        
+        return image
+        
+def get_ms(t1,t0):
+    return (t1-t0)*1000.0                  
+    
+     
+def get_largest_contours(contours):
+    areas = [cv2.contourArea(x) for x in contours]
+    max_area = max(areas)
+    max_id = areas.index(max_area)
+
+    return max_id
+    
+if __name__=='__main__':
+    image_url = '/home/thsw2/WJ/data/THexit/val/images/DJI_0645.JPG'
+    nclass = 2
+    #weights = '../weights/segmentation/BiSeNet/checkpoint.pth'
+    weights = '../weights/BiSeNet/checkpoint.pth'
+
+    segmodel = SegModel(nclass=nclass,weights=weights)
+
+    image_urls=glob.glob('/home/thsw/WJ/src/STDC/images/examples/*')
+    out_dir ='/home/thsw/WJ/src/STDC/images/resultsBise';
+    os.makedirs(out_dir,exist_ok=True)
+    for im,image_url in enumerate(image_urls[0:]):
+        #image_url = '/home/thsw2/WJ/data/THexit/val/images/54(199).JPG'
+        image_array0 = cv2.imread(image_url)     
+        time_1=time.time()        
+        pred,outstr = segmodel.eval(image_array0 )
+
+        #plt.figure(1);plt.imshow(pred);
+        #plt.show()
+        binary0 = pred.copy()
+
+        
+        time0 = time.time()
+        contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
+        max_id = -1
+        if len(contours)>0:
+            max_id = get_largest_contours(contours)
+            binary0[:,:] = 0
+            cv2.fillPoly(binary0, [contours[max_id][:,0,:]], 1) 
+      
+        time1 = time.time()
+        
+  
+        time2 = time.time()
+        
+        cv2.drawContours(image_array0,contours,max_id,(0,255,255),3)
+        time3 = time.time()
+        out_url='%s/%s'%(out_dir,os.path.basename(image_url))
+        ret = cv2.imwrite(out_url,image_array0)
+        time4 = time.time()
+        
+        print('image:%d ,infer:%.1f ms,findcontours:%.1f ms,draw:%.1f total:%.1f'%(im,get_ms(time0,time_1),get_ms(time1,time0), get_ms(time3,time2),get_ms(time3,time_1)) )
+        #print(outstr)
+        #plt.figure(0);plt.imshow(pred)
+        #plt.figure(1);plt.imshow(image_array0)
+        #plt.figure(2);plt.imshow(binary0)
+        #plt.show()
+        
+        #print(out_url,ret)
+
+
+    
+    
+
+
+
+   
+

segutils/segmodel_STDC.py

@@ -0,0 +1,128 @@
+import torch
+from torchvision import transforms
+import cv2,os,sys
+sys.path.extend(['segutils'])
+import numpy as np
+import matplotlib.pyplot as plt
+from model_stages import BiSeNet
+import torch.nn.functional as F
+import time
+
+
+class SegModel(object):
+    def __init__(self, nclass=2,weights=None,modelsize=512,device='cuda:0'):             
+
+        self.model = BiSeNet(backbone='STDCNet813', n_classes=nclass,
+                                use_boundary_2=False, use_boundary_4=False,
+                                use_boundary_8=True, use_boundary_16=False,
+                                use_conv_last=False)
+        self.device = device      
+        self.model.load_state_dict(torch.load(weights, map_location=torch.device(self.device)  ))
+        self.model= self.model.to(self.device)
+        self.mean = (0.485, 0.456, 0.406)
+        self.std = (0.229, 0.224, 0.225)
+
+    def eval(self, image=None):
+        time0 = time.time()
+        imageH, imageW, _ = image.shape
+        image = self.RB_convert(image)
+        img = self.preprocess_image(image)
+        if self.device != 'cpu':
+            imgs = img.to(self.device)
+        else:imgs=img    
+        time1 = time.time()
+        self.model.eval()
+        with torch.no_grad():
+            output = self.model(imgs)
+            
+        time2 = time.time()  
+        pred = output.data.cpu().numpy()
+        pred = np.argmax(pred, axis=1)[0]#得到每行
+        time3 = time.time() 
+        pred = cv2.resize(pred.astype(np.uint8),(imageW,imageH))
+        time4 = time.time()
+        outstr= 'pre-precess:%.1f ,infer:%.1f ,post-cpu-argmax:%.1f ,post-resize:%.1f, total:%.1f \n '%( self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3),self.get_ms(time4,time0) )
+       
+
+        return pred, outstr
+
+    def preprocess_image(self, image):
+        image = cv2.resize(image, (640,360), interpolation=cv2.INTER_LINEAR)
+        image = image.astype(np.float32)
+        image /= 255.0
+
+        image[:, :, 0] -= self.mean[0]
+        image[:, :, 1] -= self.mean[1]
+        image[:, :, 2] -= self.mean[2]
+
+        image[:, :, 0] /= self.std[0]
+        image[:, :, 1] /= self.std[1]
+        image[:, :, 2] /= self.std[2]
+
+        image = np.transpose(image, (2, 0, 1))
+        image = torch.from_numpy(image).float()
+        image = image.unsqueeze(0)
+
+        return image
+    def get_ms(self,t1,t0):
+        return (t1-t0)*1000.0     
+    def RB_convert(self,image):
+        image_c = image.copy()
+        image_c[:,:,0] = image[:,:,2]
+        image_c[:,:,2] = image[:,:,0]
+        return image_c
+
+def get_ms(t1,t0):
+    return (t1-t0)*1000.0
+def get_largest_contours(contours):
+    areas = [cv2.contourArea(x) for x in contours]
+    max_area = max(areas)
+    max_id = areas.index(max_area)
+    return max_id
+
+if __name__=='__main__':
+    impth = '/home/thsw/WJ/src/STDC/images/examples'
+    outpth= '/home/thsw/WJ/src/STDC/images/results'
+    folders = os.listdir(impth)
+    weights = '/home/thsw/WJ/src/STDC/model_maxmIOU75_1720_0.946_360640.pth'
+    segmodel = SegModel(nclass=2,weights=weights)
+    
+    for i in range(len(folders)):  
+       
+        imgpath = os.path.join(impth, folders[i])
+        time0 = time.time()
+        
+        #img = Image.open(imgpath).convert('RGB')
+        img = cv2.imread(imgpath)
+        img = np.array(img)
+        
+        
+        time1 = time.time()
+        pred, outstr = segmodel.eval(image=img)#####
+        time2 = time.time()
+        
+        binary0 = pred.copy()
+        contours, hierarchy = cv2.findContours(binary0,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
+        time3 = time.time()
+        
+        max_id = -1
+        if len(contours)>0:
+            max_id = get_largest_contours(contours)
+            binary0[:,:] = 0
+            cv2.fillPoly(binary0, [contours[max_id][:,0,:]], 1) 
+            
+        cv2.drawContours(img,contours,max_id,(0,255,255),3)
+        time4 = time.time()
+        
+        
+        #img_n = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
+        cv2.imwrite(  os.path.join( outpth,folders[i] ) ,img )
+        time5 = time.time()
+        print('image:%d ,infer:%.1f ms,findcontours:%.1f ms, draw:%.1f, total:%.1f'%(i,get_ms(time2,time1),get_ms(time3,time2),get_ms(time4,time3),get_ms(time4,time1)))                     
+
+
+    
+    
+        
+        
+ 

segutils/stdcnet.py

@@ -0,0 +1,302 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import math
+
+class ConvX(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel=3, stride=1):
+        super(ConvX, self).__init__()
+        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel, stride=stride, padding=kernel//2, bias=False)
+        self.bn = nn.BatchNorm2d(out_planes)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        out = self.relu(self.bn(self.conv(x)))
+        return out
+
+
+class AddBottleneck(nn.Module):
+    def __init__(self, in_planes, out_planes, block_num=3, stride=1):
+        super(AddBottleneck, self).__init__()
+        assert block_num > 1, print("block number should be larger than 1.")
+        self.conv_list = nn.ModuleList()
+        self.stride = stride
+        if stride == 2:
+            self.avd_layer = nn.Sequential(
+                nn.Conv2d(out_planes//2, out_planes//2, kernel_size=3, stride=2, padding=1, groups=out_planes//2, bias=False),
+                nn.BatchNorm2d(out_planes//2),
+            )
+            self.skip = nn.Sequential(
+                nn.Conv2d(in_planes, in_planes, kernel_size=3, stride=2, padding=1, groups=in_planes, bias=False),
+                nn.BatchNorm2d(in_planes),
+                nn.Conv2d(in_planes, out_planes, kernel_size=1, bias=False),
+                nn.BatchNorm2d(out_planes),
+            )
+            stride = 1
+
+        for idx in range(block_num):
+            if idx == 0:
+                self.conv_list.append(ConvX(in_planes, out_planes//2, kernel=1))
+            elif idx == 1 and block_num == 2:
+                self.conv_list.append(ConvX(out_planes//2, out_planes//2, stride=stride))
+            elif idx == 1 and block_num > 2:
+                self.conv_list.append(ConvX(out_planes//2, out_planes//4, stride=stride))
+            elif idx < block_num - 1:
+                self.conv_list.append(ConvX(out_planes//int(math.pow(2, idx)), out_planes//int(math.pow(2, idx+1))))
+            else:
+                self.conv_list.append(ConvX(out_planes//int(math.pow(2, idx)), out_planes//int(math.pow(2, idx))))
+            
+    def forward(self, x):
+        out_list = []
+        out = x
+
+        for idx, conv in enumerate(self.conv_list):
+            if idx == 0 and self.stride == 2:
+                out = self.avd_layer(conv(out))
+            else:
+                out = conv(out)
+            out_list.append(out)
+
+        if self.stride == 2:
+            x = self.skip(x)
+
+        return torch.cat(out_list, dim=1) + x
+
+
+
+class CatBottleneck(nn.Module):
+    def __init__(self, in_planes, out_planes, block_num=3, stride=1):
+        super(CatBottleneck, self).__init__()
+        assert block_num > 1, print("block number should be larger than 1.")
+        self.conv_list = nn.ModuleList()
+        self.stride = stride
+        if stride == 2:
+            self.avd_layer = nn.Sequential(
+                nn.Conv2d(out_planes//2, out_planes//2, kernel_size=3, stride=2, padding=1, groups=out_planes//2, bias=False),
+                nn.BatchNorm2d(out_planes//2),
+            )
+            self.skip = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)
+            stride = 1
+
+        for idx in range(block_num):
+            if idx == 0:
+                self.conv_list.append(ConvX(in_planes, out_planes//2, kernel=1))
+            elif idx == 1 and block_num == 2:
+                self.conv_list.append(ConvX(out_planes//2, out_planes//2, stride=stride))
+            elif idx == 1 and block_num > 2:
+                self.conv_list.append(ConvX(out_planes//2, out_planes//4, stride=stride))
+            elif idx < block_num - 1:
+                self.conv_list.append(ConvX(out_planes//int(math.pow(2, idx)), out_planes//int(math.pow(2, idx+1))))
+            else:
+                self.conv_list.append(ConvX(out_planes//int(math.pow(2, idx)), out_planes//int(math.pow(2, idx))))
+            
+    def forward(self, x):
+        out_list = []
+        out1 = self.conv_list[0](x)
+
+        for idx, conv in enumerate(self.conv_list[1:]):
+            if idx == 0:
+                if self.stride == 2:
+                    out = conv(self.avd_layer(out1))
+                else:
+                    out = conv(out1)
+            else:
+                out = conv(out)
+            out_list.append(out)
+
+        if self.stride == 2:
+            out1 = self.skip(out1)
+        out_list.insert(0, out1)
+
+        out = torch.cat(out_list, dim=1)
+        return out
+
+#STDC2Net
+class STDCNet1446(nn.Module):
+    def __init__(self, base=64, layers=[4,5,3], block_num=4, type="cat", num_classes=1000, dropout=0.20, pretrain_model='', use_conv_last=False):
+        super(STDCNet1446, self).__init__()
+        if type == "cat":
+            block = CatBottleneck
+        elif type == "add":
+            block = AddBottleneck
+        self.use_conv_last = use_conv_last
+        self.features = self._make_layers(base, layers, block_num, block)
+        self.conv_last = ConvX(base*16, max(1024, base*16), 1, 1)
+        self.gap = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(max(1024, base*16), max(1024, base*16), bias=False)
+        self.bn = nn.BatchNorm1d(max(1024, base*16))
+        self.relu = nn.ReLU(inplace=True)
+        self.dropout = nn.Dropout(p=dropout)
+        self.linear = nn.Linear(max(1024, base*16), num_classes, bias=False)
+
+        self.x2 = nn.Sequential(self.features[:1])
+        self.x4 = nn.Sequential(self.features[1:2])
+        self.x8 = nn.Sequential(self.features[2:6])
+        self.x16 = nn.Sequential(self.features[6:11])
+        self.x32 = nn.Sequential(self.features[11:])
+
+        if pretrain_model:
+            print('use pretrain model {}'.format(pretrain_model))
+            self.init_weight(pretrain_model)
+        else:
+            self.init_params()
+
+    def init_weight(self, pretrain_model):
+        
+        state_dict = torch.load(pretrain_model)["state_dict"]
+        self_state_dict = self.state_dict()
+        for k, v in state_dict.items():
+            self_state_dict.update({k: v})
+        self.load_state_dict(self_state_dict)
+
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.001)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+
+    def _make_layers(self, base, layers, block_num, block):
+        features = []
+        features += [ConvX(3, base//2, 3, 2)]
+        features += [ConvX(base//2, base, 3, 2)]
+
+        for i, layer in enumerate(layers):
+            for j in range(layer):
+                if i == 0 and j == 0:
+                    features.append(block(base, base*4, block_num, 2))
+                elif j == 0:
+                    features.append(block(base*int(math.pow(2,i+1)), base*int(math.pow(2,i+2)), block_num, 2))
+                else:
+                    features.append(block(base*int(math.pow(2,i+2)), base*int(math.pow(2,i+2)), block_num, 1))
+
+        return nn.Sequential(*features)
+
+    def forward(self, x):
+        feat2 = self.x2(x)
+        feat4 = self.x4(feat2)
+        feat8 = self.x8(feat4)
+        feat16 = self.x16(feat8)
+        feat32 = self.x32(feat16)
+        if self.use_conv_last:
+           feat32 = self.conv_last(feat32)
+
+        return feat2, feat4, feat8, feat16, feat32
+
+    def forward_impl(self, x):
+        out = self.features(x)
+        out = self.conv_last(out).pow(2)
+        out = self.gap(out).flatten(1)
+        out = self.fc(out)
+        # out = self.bn(out)
+        out = self.relu(out)
+        # out = self.relu(self.bn(self.fc(out)))
+        out = self.dropout(out)
+        out = self.linear(out)
+        return out
+
+# STDC1Net
+class STDCNet813(nn.Module):
+    def __init__(self, base=64, layers=[2,2,2], block_num=4, type="cat", num_classes=1000, dropout=0.20, pretrain_model='', use_conv_last=False):
+        super(STDCNet813, self).__init__()
+        if type == "cat":
+            block = CatBottleneck
+        elif type == "add":
+            block = AddBottleneck
+        self.use_conv_last = use_conv_last
+        self.features = self._make_layers(base, layers, block_num, block)
+        self.conv_last = ConvX(base*16, max(1024, base*16), 1, 1)
+        self.gap = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(max(1024, base*16), max(1024, base*16), bias=False)
+        self.bn = nn.BatchNorm1d(max(1024, base*16))
+        self.relu = nn.ReLU(inplace=True)
+        self.dropout = nn.Dropout(p=dropout)
+        self.linear = nn.Linear(max(1024, base*16), num_classes, bias=False)
+
+        self.x2 = nn.Sequential(self.features[:1])
+        self.x4 = nn.Sequential(self.features[1:2])
+        self.x8 = nn.Sequential(self.features[2:4])
+        self.x16 = nn.Sequential(self.features[4:6])
+        self.x32 = nn.Sequential(self.features[6:])
+
+        if pretrain_model:
+            print('use pretrain model {}'.format(pretrain_model))
+            self.init_weight(pretrain_model)
+        else:
+            self.init_params()
+
+    def init_weight(self, pretrain_model):
+        
+        state_dict = torch.load(pretrain_model)["state_dict"]
+        self_state_dict = self.state_dict()
+        for k, v in state_dict.items():
+            self_state_dict.update({k: v})
+        self.load_state_dict(self_state_dict)
+
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.001)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+
+    def _make_layers(self, base, layers, block_num, block):
+        features = []
+        features += [ConvX(3, base//2, 3, 2)]
+        features += [ConvX(base//2, base, 3, 2)]
+
+        for i, layer in enumerate(layers):
+            for j in range(layer):
+                if i == 0 and j == 0:
+                    features.append(block(base, base*4, block_num, 2))
+                elif j == 0:
+                    features.append(block(base*int(math.pow(2,i+1)), base*int(math.pow(2,i+2)), block_num, 2))
+                else:
+                    features.append(block(base*int(math.pow(2,i+2)), base*int(math.pow(2,i+2)), block_num, 1))
+
+        return nn.Sequential(*features)
+
+    def forward(self, x):
+        feat2 = self.x2(x)
+        feat4 = self.x4(feat2)
+        feat8 = self.x8(feat4)
+        feat16 = self.x16(feat8)
+        feat32 = self.x32(feat16)
+        if self.use_conv_last:
+           feat32 = self.conv_last(feat32)
+
+        return feat2, feat4, feat8, feat16, feat32
+
+    def forward_impl(self, x):
+        out = self.features(x)
+        out = self.conv_last(out).pow(2)
+        out = self.gap(out).flatten(1)
+        out = self.fc(out)
+        # out = self.bn(out)
+        out = self.relu(out)
+        # out = self.relu(self.bn(self.fc(out)))
+        out = self.dropout(out)
+        out = self.linear(out)
+        return out
+
+if __name__ == "__main__":
+    model = STDCNet813(num_classes=1000, dropout=0.00, block_num=4)
+    model.eval()
+    x = torch.randn(1,3,224,224)
+    y = model(x)
+    torch.save(model.state_dict(), 'cat.pth')
+    print(y.size())