wangjin0928
/
AIlib2
Watch 1
Star 0
Fork 0
Code Issues 0 Pull Requests 0 Releases 0 Wiki Activity
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
92 Commits
6 Branches
Branch: master
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'master'
${ noResults }
AIlib2/obbUtils/obbmodels/ctrbox_net_bak.py

205 lines
9.1KB
Raw Permalink Blame History 

  1. import torch.nn as nn

  2. import numpy as np

  3. import torch

  4. from .model_parts import CombinationModule

  5. from . import resnet

  6. import decoder

  7. 

  8. class CTRBOX_trt(nn.Module):

  9.     def __init__(self, heads, pretrained, down_ratio, final_kernel, head_conv,test_flag=False):

  10.         super(CTRBOX_trt, self).__init__()

  11. 

  12.         # channels = [3, 64, 256, 512, 1024, 2048]

  13.         # assert down_ratio in [2, 4, 8, 16]

  14.         # self.l1 = int(np.log2(down_ratio))

  15.         # self.base_network = resnet.resnet101(pretrained=pretrained)

  16.         # self.dec_c2 = CombinationModule(512, 256, batch_norm=True)

  17.         # self.dec_c3 = CombinationModule(1024, 512, batch_norm=True)

  18.         # self.dec_c4 = CombinationModule(2048, 1024, batch_norm=True)

  19. 

  20.         #channels = [3, 64, 256, 512, 1024, 2048]

  21.         #assert down_ratio in [2, 4, 8, 16]

  22.         #self.l1 = int(np.log2(down_ratio))

  23.         #self.base_network = resnet.resnet50(pretrained=pretrained)

  24.         #self.dec_c2 = CombinationModule(512, 256, batch_norm=True)

  25.         #self.dec_c3 = CombinationModule(1024, 512, batch_norm=True)

  26.         #self.dec_c4 = CombinationModule(2048, 1024, batch_norm=True)

  27. 

  28. 

  29.         #channels = [3, 64, 64, 128, 256, 512]

  30.         #assert down_ratio in [2, 4, 8, 16]

  31.         #self.l1 = int(np.log2(down_ratio))

  32.         #self.base_network = resnet.resnet34(pretrained=pretrained)

  33.         #self.dec_c2 = CombinationModule(128, 64, batch_norm=True)

  34.         #self.dec_c3 = CombinationModule(256, 128, batch_norm=True)

  35.         #self.dec_c4 = CombinationModule(512, 256, batch_norm=True)

  36. 

  37.         self.test_flag=test_flag

  38.         channels = [3, 64, 64, 128, 256, 512]

  39.         assert down_ratio in [2, 4, 8, 16]

  40.         self.l1 = int(np.log2(down_ratio))

  41.         self.base_network = resnet.resnet18(pretrained=pretrained)

  42.         self.dec_c2 = CombinationModule(128, 64, batch_norm=True)

  43.         self.dec_c3 = CombinationModule(256, 128, batch_norm=True)

  44.         self.dec_c4 = CombinationModule(512, 256, batch_norm=True)

  45. 

  46. 

  47.         self.heads = heads

  48.         

  49.         if self.test_flag:

  50.             self.decoder = decoder.DecDecoder_test(K=100,

  51.                                  conf_thresh=0.18,

  52.                                  num_classes=15)

  53.           

  54.                                  

  55. 

  56.         for head in self.heads:

  57.             classes = self.heads[head]

  58.             if head == 'wh':

  59.                 fc = nn.Sequential(nn.Conv2d(channels[self.l1], head_conv, kernel_size=3, padding=1, bias=True),

  60.                                 #    nn.BatchNorm2d(head_conv),   # BN not used in the paper, but would help stable training

  61.                                    nn.ReLU(inplace=True),

  62.                                    nn.Conv2d(head_conv, classes, kernel_size=3, padding=1, bias=True))

  63.             else:

  64.                 fc = nn.Sequential(nn.Conv2d(channels[self.l1], head_conv, kernel_size=3, padding=1, bias=True),

  65.                                 #    nn.BatchNorm2d(head_conv),   # BN not used in the paper, but would help stable training

  66.                                    nn.ReLU(inplace=True),

  67.                                    nn.Conv2d(head_conv, classes, kernel_size=final_kernel, stride=1, padding=final_kernel // 2, bias=True))

  68.             if 'hm' in head:

  69.                 fc[-1].bias.data.fill_(-2.19)

  70.             else:

  71.                 self.fill_fc_weights(fc)

  72. 

  73.             self.__setattr__(head, fc)

  74. 

  75. 

  76.     def fill_fc_weights(self, m):

  77.         if isinstance(m, nn.Conv2d):

  78.             if m.bias is not None:

  79.                 nn.init.constant_(m.bias, 0)

  80. 

  81.     def forward(self, x):

  82.         x = self.base_network(x)

  83.         # import matplotlib.pyplot as plt

  84.         # import os

  85.         # for idx in range(x[1].shape[1]):

  86.         #     temp = x[1][0,idx,:,:]

  87.         #     temp = temp.data.cpu().numpy()

  88.         #     plt.imsave(os.path.join('dilation', '{}.png'.format(idx)), temp)

  89.         c4_combine = self.dec_c4(x[-1], x[-2])

  90.         c3_combine = self.dec_c3(c4_combine, x[-3])

  91.         c2_combine = self.dec_c2(c3_combine, x[-4])

  92. 

  93.         dec_dict = {}

  94.         for head in self.heads:

  95.             dec_dict[head] = self.__getattr__(head)(c2_combine)

  96.             if 'hm' in head or 'cls' in head:

  97.                 dec_dict[head] = torch.sigmoid(dec_dict[head])

  98.                 

  99.         predictions = self.decoder.ctdet_decode(dec_dict)  

  100.         #'hm': 'wh':'reg': 'cls_theta': 

  101.         print('###############line102#############')

  102.         return predictions, dec_dict['hm'], dec_dict['wh'], dec_dict['reg'], dec_dict['cls_theta']      

  103.         #if self.test_flag:

  104.         #    predictions = self.decoder.ctdet_decode(dec_dict)  

  105.         #    return predictions

  106.         #else:    

  107.         #    return dec_dict

  108.         

  109. class CTRBOX_pth(nn.Module):

  110.     def __init__(self, heads, pretrained, down_ratio, final_kernel, head_conv,test_flag=False):

  111.         super(CTRBOX_pth, self).__init__()

  112. 

  113.         # channels = [3, 64, 256, 512, 1024, 2048]

  114.         # assert down_ratio in [2, 4, 8, 16]

  115.         # self.l1 = int(np.log2(down_ratio))

  116.         # self.base_network = resnet.resnet101(pretrained=pretrained)

  117.         # self.dec_c2 = CombinationModule(512, 256, batch_norm=True)

  118.         # self.dec_c3 = CombinationModule(1024, 512, batch_norm=True)

  119.         # self.dec_c4 = CombinationModule(2048, 1024, batch_norm=True)

  120. 

  121.         #channels = [3, 64, 256, 512, 1024, 2048]

  122.         #assert down_ratio in [2, 4, 8, 16]

  123.         #self.l1 = int(np.log2(down_ratio))

  124.         #self.base_network = resnet.resnet50(pretrained=pretrained)

  125.         #self.dec_c2 = CombinationModule(512, 256, batch_norm=True)

  126.         #self.dec_c3 = CombinationModule(1024, 512, batch_norm=True)

  127.         #self.dec_c4 = CombinationModule(2048, 1024, batch_norm=True)

  128. 

  129. 

  130.         #channels = [3, 64, 64, 128, 256, 512]

  131.         #assert down_ratio in [2, 4, 8, 16]

  132.         #self.l1 = int(np.log2(down_ratio))

  133.         #self.base_network = resnet.resnet34(pretrained=pretrained)

  134.         #self.dec_c2 = CombinationModule(128, 64, batch_norm=True)

  135.         #self.dec_c3 = CombinationModule(256, 128, batch_norm=True)

  136.         #self.dec_c4 = CombinationModule(512, 256, batch_norm=True)

  137. 

  138.         self.test_flag=test_flag

  139.         channels = [3, 64, 64, 128, 256, 512]

  140.         assert down_ratio in [2, 4, 8, 16]

  141.         self.l1 = int(np.log2(down_ratio))

  142.         self.base_network = resnet.resnet18(pretrained=pretrained)

  143.         self.dec_c2 = CombinationModule(128, 64, batch_norm=True)

  144.         self.dec_c3 = CombinationModule(256, 128, batch_norm=True)

  145.         self.dec_c4 = CombinationModule(512, 256, batch_norm=True)

  146. 

  147. 

  148.         self.heads = heads

  149.         if self.test_flag:

  150.             self.decoder = decoder.DecDecoder_test(K=100,

  151.                                  conf_thresh=0.18,

  152.                                  num_classes=15)

  153.                                  

  154. 

  155.         for head in self.heads:

  156.             classes = self.heads[head]

  157.             if head == 'wh':

  158.                 fc = nn.Sequential(nn.Conv2d(channels[self.l1], head_conv, kernel_size=3, padding=1, bias=True),

  159.                                 #    nn.BatchNorm2d(head_conv),   # BN not used in the paper, but would help stable training

  160.                                    nn.ReLU(inplace=True),

  161.                                    nn.Conv2d(head_conv, classes, kernel_size=3, padding=1, bias=True))

  162.             else:

  163.                 fc = nn.Sequential(nn.Conv2d(channels[self.l1], head_conv, kernel_size=3, padding=1, bias=True),

  164.                                 #    nn.BatchNorm2d(head_conv),   # BN not used in the paper, but would help stable training

  165.                                    nn.ReLU(inplace=True),

  166.                                    nn.Conv2d(head_conv, classes, kernel_size=final_kernel, stride=1, padding=final_kernel // 2, bias=True))

  167.             if 'hm' in head:

  168.                 fc[-1].bias.data.fill_(-2.19)

  169.             else:

  170.                 self.fill_fc_weights(fc)

  171. 

  172.             self.__setattr__(head, fc)

  173. 

  174. 

  175.     def fill_fc_weights(self, m):

  176.         if isinstance(m, nn.Conv2d):

  177.             if m.bias is not None:

  178.                 nn.init.constant_(m.bias, 0)

  179. 

  180.     def forward(self, x):

  181.         x = self.base_network(x)

  182.         # import matplotlib.pyplot as plt

  183.         # import os

  184.         # for idx in range(x[1].shape[1]):

  185.         #     temp = x[1][0,idx,:,:]

  186.         #     temp = temp.data.cpu().numpy()

  187.         #     plt.imsave(os.path.join('dilation', '{}.png'.format(idx)), temp)

  188.         c4_combine = self.dec_c4(x[-1], x[-2])

  189.         c3_combine = self.dec_c3(c4_combine, x[-3])

  190.         c2_combine = self.dec_c2(c3_combine, x[-4])

  191. 

  192.         dec_dict = {}

  193.         for head in self.heads:

  194.             dec_dict[head] = self.__getattr__(head)(c2_combine)

  195.             if 'hm' in head or 'cls' in head:

  196.                 dec_dict[head] = torch.sigmoid(dec_dict[head])

  197.   

  198.         

  199.         if self.test_flag:

  200.             predictions = self.decoder.ctdet_decode(dec_dict)  

  201.             print('##line301:',predictions )

  202.             return predictions, dec_dict['hm'], dec_dict['wh'], dec_dict['reg'], dec_dict['cls_theta']  

  203.         else:    

  204.             return dec_dict