把代码整理了一下 发出来
#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)本文章为转载内容，我们尊重原作者对文章享有的著作权。如有内容错误或侵权问题，欢迎原作者联系我们进行内容更正或删除文章。
