用pytorch的VGG网络训练自己的数据集，以及如何加载预训练网络

最新推荐文章于 2024-07-23 10:06:30 发布

若曦爹

最新推荐文章于 2024-07-23 10:06:30 发布

阅读量3.1k

点赞数 3

本文链接：https://blog.csdn.net/gbz3300255/article/details/107614203

版权

把代码整理了一下发出来

#encoding=utf-8
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim as optim  
import matplotlib.pyplot as plt
import numpy as np


def loadtraindata():
    path = r"/home/yuany/classificationmaster/data/train"                                         # 路径
    trainset = torchvision.datasets.ImageFolder(path,
                                                transform=transforms.Compose([
                                                    transforms.Resize((224, 224)),  # 将图片缩放到指定大小（h,w）或者保持长宽比并缩放最短的边到int大小

                                                    transforms.CenterCrop(224),
                                                    transforms.ToTensor()])
                                                )

    trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                              shuffle=True, num_workers=2)
    return trainloader




def Conv3x3BNReLU(in_channels,out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels=in_channels,out_channels=out_channels,kernel_size=3,stride=1,padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU6(inplace=True)
    )

class VGGNet(nn.Module):
    def __init__(self, block_nums,num_classes=7):
        super(VGGNet, self).__init__()

        self.stage1 = self._make_layers(in_channels=3, out_channels=64, block_num=block_nums[0])
        self.stage2 = self._make_layers(in_channels=64, out_channels=128, block_num=block_nums[1])
        self.stage3 = self._make_layers(in_channels=128, out_channels=256, block_num=block_nums[2])
        self.stage4 = self._make_layers(in_channels=256, out_channels=512, block_num=block_nums[3])
        self.stage5 = self._make_layers(in_channels=512, out_channels=512, block_num=block_nums[4])

        self.classifier = nn.Sequential(
            nn.Linear(in_features=512*7*7,out_features=4096),
            nn.ReLU6(inplace=True),
            nn.Dropout(p=0.2),
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU6(inplace=True),
            nn.Dropout(p=0.2),
            nn.Linear(in_features=4096, out_features=num_classes)
        )

    def _make_layers(self, in_channels, out_channels, block_num):
        layers = []
        layers.append(Conv3x3BNReLU(in_channels,out_channels))
        for i in range(1,block_num):
            layers.append(Conv3x3BNReLU(out_channels,out_channels))
        layers.append(nn.MaxPool2d(kernel_size=2,stride=2))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = self.stage5(x)
        x = x.view(x.size(0),-1)
        out = self.classifier(x)
        return out

def VGG16():
    block_nums = [2, 2, 3, 3, 3]
    model = VGGNet(block_nums)
    return model

def VGG19():
    block_nums = [2, 2, 4, 4, 4]
    model = VGGNet(block_nums)
    return model

	
	

class Net(nn.Module):                 # 定义网络，继承torch.nn.Module
	def __init__(self):
		super(Net, self).__init__()   
		self.conv1 = nn.Conv2d(3, 6, 5)       # 卷积层
		self.pool = nn.MaxPool2d(2, 2)        # 池化层
		self.conv2 = nn.Conv2d(6, 16, 5)      # 卷积层
		self.fc1 = nn.Linear(16 * 5 * 5, 120) # 全连接层
		self.fc2 = nn.Linear(120, 84)
		self.fc3 = nn.Linear(84, 7)      # 10个输出
	
	
	def forward(self, x):                   # 前向传播                                   
		x = self.pool(F.relu(self.conv1(x)))  # F就是torch.nn.functional
		x = self.pool(F.relu(self.conv2(x)))
		x = x.view(-1, 16 * 5 * 5)  # .view( )是一个tensor的方法，使得tensor改变size但是元素的总数是不变的。从卷基层到全连接层的维度转换
		
		x = F.relu(self.fc1(x))
		x = F.relu(self.fc2(x))
		x = self.fc3(x)
		return x
classes = ('0','1', '2', '3', '4','5', '6')

def loadtestdata():
    path = r"/home/yuany/classificationmaster/data/val"
    testset = torchvision.datasets.ImageFolder(path,
                                                transform=transforms.Compose([
                                                    transforms.Resize((32, 32)),  # 将图片缩放到指定大小（h,w）或者保持长宽比并缩放最短的边到int大小
                                                    transforms.ToTensor()])
                                                )
    testloader = torch.utils.data.DataLoader(testset, batch_size=1,
                                             shuffle=True, num_workers=2)
    return testloader



def trainandsave():
    trainloader = loadtraindata()
    # 神经网络结构
    #net = Net()
    net = VGG16()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("device = ", device)
    net.to(device=device)
    
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)   # 学习率为0.001
    criterion = nn.CrossEntropyLoss()   # 损失函数也可以自己定义，我们这里用的交叉熵损失函数
    # 训练部分
    print("trainloader = ", trainloader)
    for epoch in range(20):    # 训练的数据量为5个epoch，每个epoch为一个循环
                            # 每个epoch要训练所有的图片，每训练完成200张便打印一下训练的效果（loss值）
        running_loss = 0.0  # 定义一个变量方便我们对loss进行输出
        for i, data in enumerate(trainloader, 0):  # 这里我们遇到了第一步中出现的trailoader，代码传入数据
            # enumerate是python的内置函数，既获得索引也获得数据
            # get the inputs
            inputs, labels = data  # data是从enumerate返回的data，包含数据和标签信息，分别赋值给inputs和labels
            
            #print("inputs = ", inputs)
            #print("labels = ", labels)
            # wrap them in Variable
            inputs, labels = Variable(inputs), Variable(labels)  # 转换数据格式用Variable
           
            optimizer.zero_grad()        # 梯度置零，因为反向传播过程中梯度会累加上一次循环的梯度

            # forward + backward + optimize
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = net(inputs)        # 把数据输进CNN网络net
            
            loss = criterion(outputs, labels)  # 计算损失值
            loss.backward()                    # loss反向传播 计算反向梯度
            optimizer.step()                   # 利用反向梯度 参数更新 
            #running_loss += loss.data[0]       # loss累加
            running_loss += loss.item()       # loss累加
            
            if i % 2 == 1:                 
                print('[%d, %5d] loss: %.3f' %
                    (epoch + 1, i + 1, running_loss / 2))  # 然后再除以200，就得到这两百次的平均损失值
                running_loss = 0.0  # 这一个200次结束后，就把running_loss归零，下一个200次继续使用

    print('Finished Training')
    # 保存神经网络
    torch.save(net, 'net.pkl')                      # 保存整个神经网络的结构和模型参数
    torch.save(net.state_dict(), 'net_params.pkl')  # 只保存神经网络的模型参数



def reload_net():
    trainednet = torch.load('net.pkl')
    return trainednet

def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def test():
	testloader = loadtestdata()
	net = reload_net()
	
	device = torch.device('cuda')
	net.to(device=device)
	
	#dataiter = iter(testloader)  
	total = 0
	right = 0
	for i, data in enumerate(testloader, 0):  # 这里我们遇到了第一步中出现的trailoader，代码传入数据
		inputs, labels = data  # data是从enumerate返回的data，包含数据和标签信息，分别赋值给inputs和labels
		
		inputs = Variable(inputs)
		inputs = inputs.to(device)
		labels = labels.to(device)
		
		outputs = net(inputs)  

		_, predicted = torch.max(outputs.data, 1)
		
		out = F.softmax(outputs, dim=1)                    # 采用SoftMax方法将输出的2个输出值调整至[0.0, 1.0],两者和为1
	
		print(out) 
		print("predicted = ", predicted) 
		print("index = ", out.argmax(dim=1))
		print("maxvalue = ", torch.max(out))
		
		b = (torch.max(out)).cpu()
		c = b.detach().numpy()
		print("c = ", c)
	
		total = total + 1
		if predicted == labels:
			right = right + 1
		
		print("per = ", total / 19084)

	print("total = ", total)
	print("right = ", right)
	print("ratio = ", right / total)
			
		

if __name__ == '__main__':
	trainandsave()
	#test()


opt_SGD=torch.optim.SGD(net_SGD.parameters(),lr=LR)
opt_Momentum=torch.optim.SGD(net_Momentum.parameters(),lr=LR,momentum=0.8)
opt_RMSprop=torch.optim.RMSprop(net_RMSprop.parameters(),lr=LR,alpha=0.9)
opt_Adam=torch.optim.Adam(net_Adam.parameters(),lr=LR,betas=(0.9,0.99)) 
四种优化器 自己选择着用吧

训练代码由个很重要的问题，如何加载预训练网络。pytorch提供了两种方法，但是 save和load要配套使用才好哦

保存和加载模型参数有两种方式：

方式一：

torch.save(net.state_dict(),path):

功能：保存训练完的网络的各层参数（即weights和bias)

其中：net.state_dict()获取各层参数，path是文件存放路径(通常保存文件格式为.pt或.pth)

net2.load_state_dict(torch.load(path)):

功能：加载保存到path中的各层参数到神经网络

注意：不可以直接为torch.load_state_dict(path)，此函数不能直接接收字符串类型参数

方式二：

torch.save(net,path):

功能：保存训练完的整个网络模型（不止weights和bias）

net2=torch.load(path):

功能：加载保存到path中的整个神经网络

说明：官方推荐方式一，原因自然是保存的内容少，速度会更快。

切记切记

追加了怎么获得图像类别置信度的操作盗图来也

多分类神经网络中，最后一层通常是linear+softmax，为什么换成linear+relu+softmax后，正确率就急剧下降？

答：你把负值全给扔了，事实上在训练中，一开始网络的正确率很低；往往会把一些正确的类分配为负值，然后relu把他们都变成0。。。。训练效果肯定上不去

下面这个代码包含的内容太多了。例如pytorch如何标记连通区域。连通区域标出来怎么用。如何将数据存成txt，怎么将txt加载进来读成矩阵形式。等等吧

#encoding=utf-8
import cv2
import numpy as np
import torch
from skimage import measure

import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim as optim  
import matplotlib.pyplot as plt
from PIL import Image

import warnings
from collections import namedtuple
from torch.utils.model_zoo import load_url as load_state_dict_from_url

###################################################### VGGNET #########################################################
def Conv3x3BNReLU(in_channels,out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels=in_channels,out_channels=out_channels,kernel_size=3,stride=1,padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU6(inplace=True)
    )

class VGGNet(nn.Module):
    def __init__(self, block_nums,num_classes=7):
        super(VGGNet, self).__init__()

        self.stage1 = self._make_layers(in_channels=3, out_channels=64, block_num=block_nums[0])
        self.stage2 = self._make_layers(in_channels=64, out_channels=128, block_num=block_nums[1])
        self.stage3 = self._make_layers(in_channels=128, out_channels=256, block_num=block_nums[2])
        self.stage4 = self._make_layers(in_channels=256, out_channels=512, block_num=block_nums[3])
        self.stage5 = self._make_layers(in_channels=512, out_channels=512, block_num=block_nums[4])

        self.classifier = nn.Sequential(
            nn.Linear(in_features=512*7*7,out_features=4096),
            nn.ReLU6(inplace=True),
            nn.Dropout(p=0.2),
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU6(inplace=True),
            nn.Dropout(p=0.2),
            nn.Linear(in_features=4096, out_features=num_classes)
        )

    def _make_layers(self, in_channels, out_channels, block_num):
        layers = []
        layers.append(Conv3x3BNReLU(in_channels,out_channels))
        for i in range(1,block_num):
            layers.append(Conv3x3BNReLU(out_channels,out_channels))
        layers.append(nn.MaxPool2d(kernel_size=2,stride=2))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = self.stage5(x)
        x = x.view(x.size(0),-1)
        out = self.classifier(x)
        return out

def VGG16():
    block_nums = [2, 2, 3, 3, 3]
    model = VGGNet(block_nums)
    return model

def VGG19():
    block_nums = [2, 2, 4, 4, 4]
    model = VGGNet(block_nums)
    return model
	
classes = ('0','1', '2', '3', '4','5','6')

###################################################### MobileNetV2 #########################################################
class LinearBottleNeck(nn.Module):
    def __init__(self, in_channels, out_channels, stride, t=6, class_num=7):
        super().__init__()

        self.residual = nn.Sequential(
            nn.Conv2d(in_channels, in_channels * t, 1),
            nn.BatchNorm2d(in_channels * t),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels * t, in_channels * t, 3, stride=stride, padding=1, groups=in_channels * t),
            nn.BatchNorm2d(in_channels * t),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels * t, out_channels, 1),
            nn.BatchNorm2d(out_channels)
        )

        self.stride = stride
        self.in_channels = in_channels
        self.out_channels = out_channels
    
    def forward(self, x):
        residual = self.residual(x)

        if self.stride == 1 and self.in_channels == self.out_channels:
            residual += x
        
        return residual

class MobileNetV2(nn.Module):
    def __init__(self, class_num=7):
        super().__init__()

        self.pre = nn.Sequential(
            nn.Conv2d(3, 32, 1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU6(inplace=True)
        )

        self.stage1 = LinearBottleNeck(32, 16, 1, 1)
        self.stage2 = self._make_stage(2, 16, 24, 2, 6)
        self.stage3 = self._make_stage(3, 24, 32, 2, 6)
        self.stage4 = self._make_stage(4, 32, 64, 2, 6)
        self.stage5 = self._make_stage(3, 64, 96, 1, 6)
        self.stage6 = self._make_stage(3, 96, 160, 1, 6)
        self.stage7 = LinearBottleNeck(160, 320, 1, 6)

        self.conv1 = nn.Sequential(
            nn.Conv2d(320, 1280, 1),
            nn.BatchNorm2d(1280),
            nn.ReLU6(inplace=True)
        )

        self.conv2 = nn.Conv2d(1280, class_num, 1)
            
    def forward(self, x):
        x = self.pre(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = self.stage5(x)
        x = self.stage6(x)
        x = self.stage7(x)
        x = self.conv1(x)
        x = F.adaptive_avg_pool2d(x, 1)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)

        return x
    
    def _make_stage(self, repeat, in_channels, out_channels, stride, t):

        layers = []
        layers.append(LinearBottleNeck(in_channels, out_channels, stride, t))
        
        while repeat - 1:
            layers.append(LinearBottleNeck(out_channels, out_channels, 1, t))
            repeat -= 1
        
        return nn.Sequential(*layers)
        
        
def mobilenetv2():
    return MobileNetV2() 

###################################################### googlenet #########################################################
__all__ = ['GoogLeNet', 'googlenet']

model_urls = {
    # GoogLeNet ported from TensorFlow
    'googlenet': 'https://download.pytorch.org/models/googlenet-1378be20.pth',
}

_GoogLeNetOuputs = namedtuple('GoogLeNetOuputs', ['logits', 'aux_logits2', 'aux_logits1'])


def googlenet(pretrained=False, progress=True, **kwargs):
    r"""GoogLeNet (Inception v1) model architecture from
    `"Going Deeper with Convolutions" <http://arxiv.org/abs/1409.4842>`_.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
        aux_logits (bool): If True, adds two auxiliary branches that can improve training.
            Default: *False* when pretrained is True otherwise *True*
        transform_input (bool): If True, preprocesses the input according to the method with which it
            was trained on ImageNet. Default: *False*
    """
    if pretrained:
        if 'transform_input' not in kwargs:
            kwargs['transform_input'] = True
        if 'aux_logits' not in kwargs:
            kwargs['aux_logits'] = False
        if kwargs['aux_logits']:
            warnings.warn('auxiliary heads in the pretrained googlenet model are NOT pretrained, '
                          'so make sure to train them')
        original_aux_logits = kwargs['aux_logits']
        kwargs['aux_logits'] = True
        kwargs['init_weights'] = False
        model = GoogLeNet(**kwargs)
        state_dict = load_state_dict_from_url(model_urls['googlenet'],
                                              progress=progress)
        model.load_state_dict(state_dict)
        if not original_aux_logits:
            model.aux_logits = False
            del model.aux1, model.aux2
        return model

    return GoogLeNet(**kwargs)

class GoogLeNet(nn.Module):

    def __init__(self, num_classes=7, aux_logits=True, transform_input=False, init_weights=True):
        super(GoogLeNet, self).__init__()
        self.aux_logits = aux_logits
        self.transform_input = transform_input

        self.conv1 = BasicConv2d(3, 64, kernel_size=7, stride=2, padding=3)
        self.maxpool1 = nn.MaxPool2d(3, stride=2, ceil_mode=True)
        self.conv2 = BasicConv2d(64, 64, kernel_size=1)
        self.conv3 = BasicConv2d(64, 192, kernel_size=3, padding=1)
        self.maxpool2 = nn.MaxPool2d(3, stride=2, ceil_mode=True)

        self.inception3a = Inception(192, 64, 96, 128, 16, 32, 32)
        self.inception3b = Inception(256, 128, 128, 192, 32, 96, 64)
        self.maxpool3 = nn.MaxPool2d(3, stride=2, ceil_mode=True)

        self.inception4a = Inception(480, 192, 96, 208, 16, 48, 64)
        self.inception4b = Inception(512, 160, 112, 224, 24, 64, 64)
        self.inception4c = Inception(512, 128, 128, 256, 24, 64, 64)
        self.inception4d = Inception(512, 112, 144, 288, 32, 64, 64)
        self.inception4e = Inception(528, 256, 160, 320, 32, 128, 128)
        self.maxpool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)

        self.inception5a = Inception(832, 256, 160, 320, 32, 128, 128)
        self.inception5b = Inception(832, 384, 192, 384, 48, 128, 128)

        if aux_logits:
            self.aux1 = InceptionAux(512, num_classes)
            self.aux2 = InceptionAux(528, num_classes)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(0.2)
        self.fc = nn.Linear(1024, num_classes)

        if init_weights:
            self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                import scipy.stats as stats
                X = stats.truncnorm(-2, 2, scale=0.01)
                values = torch.as_tensor(X.rvs(m.weight.numel()), dtype=m.weight.dtype)
                values = values.view(m.weight.size())
                with torch.no_grad():
                    m.weight.copy_(values)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        if self.transform_input:
            x_ch0 = torch.unsqueeze(x[:, 0], 1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
            x_ch1 = torch.unsqueeze(x[:, 1], 1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
            x_ch2 = torch.unsqueeze(x[:, 2], 1) * (0.225 / 0.5) + (0.406 - 0.5) / 0.5
            x = torch.cat((x_ch0, x_ch1, x_ch2), 1)

        # N x 3 x 224 x 224
        x = self.conv1(x)
        # N x 64 x 112 x 112
        x = self.maxpool1(x)
        # N x 64 x 56 x 56
        x = self.conv2(x)
        # N x 64 x 56 x 56
        x = self.conv3(x)
        # N x 192 x 56 x 56
        x = self.maxpool2(x)

        # N x 192 x 28 x 28
        x = self.inception3a(x)
        # N x 256 x 28 x 28
        x = self.inception3b(x)
        # N x 480 x 28 x 28
        x = self.maxpool3(x)
        # N x 480 x 14 x 14
        x = self.inception4a(x)
        # N x 512 x 14 x 14
        if self.training and self.aux_logits:
            aux1 = self.aux1(x)

        x = self.inception4b(x)
        # N x 512 x 14 x 14
        x = self.inception4c(x)
        # N x 512 x 14 x 14
        x = self.inception4d(x)
        # N x 528 x 14 x 14
        if self.training and self.aux_logits:
            aux2 = self.aux2(x)

        x = self.inception4e(x)
        # N x 832 x 14 x 14
        x = self.maxpool4(x)
        # N x 832 x 7 x 7
        x = self.inception5a(x)
        # N x 832 x 7 x 7
        x = self.inception5b(x)
        # N x 1024 x 7 x 7

        x = self.avgpool(x)
        # N x 1024 x 1 x 1
        x = x.view(x.size(0), -1)
        # N x 1024
        x = self.dropout(x)
        x = self.fc(x)
        # N x 1000 (num_classes)
        if self.training and self.aux_logits:
            return _GoogLeNetOuputs(x, aux2, aux1)
        return x

class Inception(nn.Module):

    def __init__(self, in_channels, ch1x1, ch3x3red, ch3x3, ch5x5red, ch5x5, pool_proj):
        super(Inception, self).__init__()

        self.branch1 = BasicConv2d(in_channels, ch1x1, kernel_size=1)

        self.branch2 = nn.Sequential(
            BasicConv2d(in_channels, ch3x3red, kernel_size=1),
            BasicConv2d(ch3x3red, ch3x3, kernel_size=3, padding=1)
        )

        self.branch3 = nn.Sequential(
            BasicConv2d(in_channels, ch5x5red, kernel_size=1),
            BasicConv2d(ch5x5red, ch5x5, kernel_size=3, padding=1)
        )

        self.branch4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1, ceil_mode=True),
            BasicConv2d(in_channels, pool_proj, kernel_size=1)
        )

    def forward(self, x):
        branch1 = self.branch1(x)
        branch2 = self.branch2(x)
        branch3 = self.branch3(x)
        branch4 = self.branch4(x)

        outputs = [branch1, branch2, branch3, branch4]
        return torch.cat(outputs, 1)


class InceptionAux(nn.Module):

    def __init__(self, in_channels, num_classes):
        super(InceptionAux, self).__init__()
        self.conv = BasicConv2d(in_channels, 128, kernel_size=1)

        self.fc1 = nn.Linear(2048, 1024)
        self.fc2 = nn.Linear(1024, num_classes)

    def forward(self, x):
        # aux1: N x 512 x 14 x 14, aux2: N x 528 x 14 x 14
        x = F.adaptive_avg_pool2d(x, (4, 4))
        # aux1: N x 512 x 4 x 4, aux2: N x 528 x 4 x 4
        x = self.conv(x)
        # N x 128 x 4 x 4
        x = x.view(x.size(0), -1)
        # N x 2048
        x = F.relu(self.fc1(x), inplace=True)
        # N x 1024
        x = F.dropout(x, 0.7, training=self.training)
        # N x 1024
        x = self.fc2(x)
        # N x num_classes

        return x


class BasicConv2d(nn.Module):       #Conv2d+BN+Relu

    def __init__(self, in_channels, out_channels, **kwargs):
        super(BasicConv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
        self.bn = nn.BatchNorm2d(out_channels, eps=0.001)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return F.relu(x, inplace=True)

def googlenetv1():
    return GoogLeNet() 


def reload_net():
    trainednet = torch.load('bestnet.pkl')
    return trainednet



def CreateTxt(Gbig, wimg, himg):
	f = "lucky.txt"
	with open(f,"a") as file:
		file.write("width = " + str(wimg) + ", "+ "height = " + str(himg) + ", "+ "channels = " + str(1) +"\n")
		file.write("0:null,1:line" +"\n")
		for u in range(himg):
			for v in range(wimg):
				gray = Gbig[u, v]
				if gray == 255:
					file.write(str(1) + " ")
				else:
					file.write(str(0) + " ")
			file.write("\n")

def ReadTxt(wimg, himg, BINARY):
	
	A = np.zeros((himg,wimg),dtype=int)
	f = open('lucky.txt')
	lines = f.readlines()
	A_row = 0
	adds = 0;
	
	'''
	for line in lines:
		if adds < 2:
			A_row = 0
		else:
			#print("line = ", line)
			list = line.strip('\n').split(' ')  
			A[A_row:] = list[0:wimg]
			A_row += 1
		adds += 1
	'''	
	a = np.loadtxt('lucky.txt', skiprows=2, dtype=int)
	print("a = ", a)		
	print("a shape = ", a.shape)	
	for u in range(himg):
		for v in range(wimg):
			gray = a[u, v]
			if gray == 1:
				BINARY[u, v] = 255

def threshold(image, binary, orgimg):
	#gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
	gray = image
	
	orgimgLAB = orgimg.copy()
	
	tpath2='22222.jpg'
	cv2.imwrite(tpath2, gray)
	
	ret, binary =  cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
	himg, wimg = gray.shape
	
	#CreateTxt(binary, wimg, himg)

	BINARY = np.zeros((himg,wimg),dtype=np.uint8)
	ReadTxt(wimg, himg, BINARY)
	
	tpathx='333333.jpg'
	cv2.imwrite(tpathx, BINARY)

	imgxx = image.copy()
	lable =  measure.label(BINARY,  connectivity = 2)

	props = measure.regionprops(lable, intensity_image=None, cache=True)
	
	net = reload_net()
	device = torch.device('cuda')
	net.to(device=device)
	
	for i in range(len(props)):
		#print(props[i].label)
		#print(props[i].bbox)
		y0, x0, y1, x1 = props[i].bbox
		if (x1 - x0) > 60 or (y1 - y0) > 60 or ((x1 - x0) < 30 and (y1 - y0) < 30):
			cv2.rectangle(imgxx, (x0, y0), (x1, y1), (100, 100, 0), 1)
		else:
			xcenter = (x0 + x1) / 2
			ycenter = (y0 + y1) / 2
			warrow = x1 - x0
			harrow = y1 - y0
			newx00 = 0
			newx11 = 0
			newy00 = 0
			newx11 = 0
			if warrow >=  harrow:
				newx00 = x0
				newx11 = x1
				newy00 = int(ycenter - (warrow) / 2)
				newy11 = int(ycenter + (warrow) / 2)
			else:
				newy00 = y0
				newy11 = y1

				newx00 = int(xcenter - (harrow) / 2)
				newx11 = int(xcenter + (harrow) / 2)	
			cropImg1 = orgimg[newy00:newy11, newx00:newx11]
			
			cropImg = cropImg1.copy()
			
			chimg,cwimg,ccimg = cropImg.shape
			
			r = chimg * chimg / 4
			centerx = chimg / 2
			centery = cwimg / 2
			
			
			for u in range(chimg):
				for v in range(cwimg):
					rad = (u - centery) * (u - centery) + (v - centerx) * (v - centerx)
					if rad > r:
						cropImg[u, v, 0] = 88
						cropImg[u, v, 1] = 91
						cropImg[u, v, 2] = 96
			
			img_pathx = str(i + 200000) + ".jpg"
			tpath = "./outdata/" + 	img_pathx
			cv2.imwrite(tpath, cropImg)
			
			image = Image.fromarray(cv2.cvtColor(cropImg,cv2.COLOR_BGR2RGB))
			
			
			transform=transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor()])
			tranto = transform(image)
			
			tranto=tranto[np.newaxis,:, : ,:]
			#print("tranto shape = ", tranto.shape)
			
			inputs = Variable(tranto)
			#print("inputs shape = ", inputs.shape)
			inputs = inputs.to(device)
			#outputs,x1,x2 = net(inputs)  
			outputs = net(inputs)  
			_, predicted = torch.max(outputs.data, 1)
			out = F.softmax(outputs, dim=1)
			b = (torch.max(out)).cpu()
			confidence = b.detach().numpy()
			lablepre = (out.argmax(dim=1))
			#lablepre = b.detach().numpy()
			#print("confidence = ", confidence)
			lablepre = predicted.cpu().detach().numpy()
			#print("lablepre = ", lablepre)
			if lablepre == [0]:
				cv2.rectangle(orgimgLAB, (x0, y0), (x1, y1), (0, 0, 255), 1) 
			elif lablepre == [1]:
				cv2.rectangle(orgimgLAB, (x0, y0), (x1, y1), (0, 255, 0), 1) 
			elif lablepre == [2]:
				cv2.rectangle(orgimgLAB, (x0, y0), (x1, y1), (255, 0, 0), 1) 
			elif lablepre == [3]:
				cv2.rectangle(orgimgLAB, (x0, y0), (x1, y1), (255, 255, 0), 1)
			elif lablepre == [4]:
				cv2.rectangle(orgimgLAB, (x0, y0), (x1, y1), (0, 255, 255), 1) 
			elif lablepre == [5]:
				cv2.rectangle(orgimgLAB, (x0, y0), (x1, y1), (255, 0, 255), 1)
			elif lablepre == [6]:
				cv2.rectangle(orgimgLAB, (x0, y0), (x1, y1), (255, 255, 255), 1)
			else:
				print("error")
	
	tpath1='1111.jpg'
	cv2.imwrite(tpath1, orgimgLAB)
		

if __name__ == '__main__':
	src = cv2.imread('input.jpg')
	Bbig,Gbig,Rbig = cv2.split(src)
	himg,wimg,cimg = src.shape
	oimg = cv2.imread('./data/test/input/10_03.jpg')
	threshold(Gbig, Bbig, oimg)