first push
This commit is contained in:
BIN
utils/.DS_Store
vendored
Normal file
BIN
utils/.DS_Store
vendored
Normal file
Binary file not shown.
4
utils/celeryconfig.py
Normal file
4
utils/celeryconfig.py
Normal file
@ -0,0 +1,4 @@
|
||||
broker_url = 'redis://127.0.0.1:6379/0'
|
||||
result_backend = 'redis://127.0.0.1:6379/1'
|
||||
includ = ['ieemoo-ai-search.ieemoo-ai-search']
|
||||
main = 'ieemoo-ai-search.ieemoo-ai-search'
|
||||
582
utils/classify.py
Normal file
582
utils/classify.py
Normal file
@ -0,0 +1,582 @@
|
||||
# coding=utf-8
|
||||
# /usr/bin/env pythpn
|
||||
|
||||
'''
|
||||
Author: yinhao
|
||||
Email: yinhao_x@163.com
|
||||
Wechat: xss_yinhao
|
||||
Github: http://github.com/yinhaoxs
|
||||
data: 2019-11-23 18:29
|
||||
desc:
|
||||
'''
|
||||
|
||||
import torch.nn as nn
|
||||
import math
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
from torch import nn
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.autograd import Variable
|
||||
import cv2
|
||||
import shutil
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
import os
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
|
||||
# config.py
|
||||
BATCH_SIZE = 16
|
||||
PROPOSAL_NUM = 6
|
||||
CAT_NUM = 4
|
||||
INPUT_SIZE = (448, 448) # (w, h)
|
||||
DROP_OUT = 0.5
|
||||
CLASS_NUM = 37
|
||||
|
||||
|
||||
# resnet.py
|
||||
def conv3x3(in_planes, out_planes, stride=1):
|
||||
"3x3 convolution with padding"
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * 4)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Module):
|
||||
def __init__(self, block, layers, num_classes=1000):
|
||||
self.inplanes = 64
|
||||
super(ResNet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
|
||||
bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
||||
self.avgpool = nn.AvgPool2d(7)
|
||||
self.fc = nn.Linear(512 * block.expansion, num_classes)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(self.inplanes, planes * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
feature1 = x
|
||||
x = self.avgpool(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
x = nn.Dropout(p=0.5)(x)
|
||||
feature2 = x
|
||||
x = self.fc(x)
|
||||
|
||||
return x, feature1, feature2
|
||||
|
||||
|
||||
# model.py
|
||||
class ProposalNet(nn.Module):
|
||||
def __init__(self):
|
||||
super(ProposalNet, self).__init__()
|
||||
self.down1 = nn.Conv2d(2048, 128, 3, 1, 1)
|
||||
self.down2 = nn.Conv2d(128, 128, 3, 2, 1)
|
||||
self.down3 = nn.Conv2d(128, 128, 3, 2, 1)
|
||||
self.ReLU = nn.ReLU()
|
||||
self.tidy1 = nn.Conv2d(128, 6, 1, 1, 0)
|
||||
self.tidy2 = nn.Conv2d(128, 6, 1, 1, 0)
|
||||
self.tidy3 = nn.Conv2d(128, 9, 1, 1, 0)
|
||||
|
||||
def forward(self, x):
|
||||
batch_size = x.size(0)
|
||||
d1 = self.ReLU(self.down1(x))
|
||||
d2 = self.ReLU(self.down2(d1))
|
||||
d3 = self.ReLU(self.down3(d2))
|
||||
t1 = self.tidy1(d1).view(batch_size, -1)
|
||||
t2 = self.tidy2(d2).view(batch_size, -1)
|
||||
t3 = self.tidy3(d3).view(batch_size, -1)
|
||||
return torch.cat((t1, t2, t3), dim=1)
|
||||
|
||||
|
||||
class AttentionNet(nn.Module):
|
||||
def __init__(self, topN=4):
|
||||
super(attention_net, self).__init__()
|
||||
self.pretrained_model = ResNet(Bottleneck, [3, 4, 6, 3])
|
||||
self.pretrained_model.avgpool = nn.AdaptiveAvgPool2d(1)
|
||||
self.pretrained_model.fc = nn.Linear(512 * 4, 200)
|
||||
self.proposal_net = ProposalNet()
|
||||
self.topN = topN
|
||||
self.concat_net = nn.Linear(2048 * (CAT_NUM + 1), 200)
|
||||
self.partcls_net = nn.Linear(512 * 4, 200)
|
||||
_, edge_anchors, _ = generate_default_anchor_maps()
|
||||
self.pad_side = 224
|
||||
self.edge_anchors = (edge_anchors + 224).astype(np.int)
|
||||
|
||||
def forward(self, x):
|
||||
resnet_out, rpn_feature, feature = self.pretrained_model(x)
|
||||
x_pad = F.pad(x, (self.pad_side, self.pad_side, self.pad_side, self.pad_side), mode='constant', value=0)
|
||||
batch = x.size(0)
|
||||
# we will reshape rpn to shape: batch * nb_anchor
|
||||
rpn_score = self.proposal_net(rpn_feature.detach())
|
||||
all_cdds = [
|
||||
np.concatenate((x.reshape(-1, 1), self.edge_anchors.copy(), np.arange(0, len(x)).reshape(-1, 1)), axis=1)
|
||||
for x in rpn_score.data.cpu().numpy()]
|
||||
top_n_cdds = [hard_nms(x, topn=self.topN, iou_thresh=0.25) for x in all_cdds]
|
||||
top_n_cdds = np.array(top_n_cdds)
|
||||
top_n_index = top_n_cdds[:, :, -1].astype(np.int)
|
||||
top_n_index = torch.from_numpy(top_n_index).cuda()
|
||||
top_n_prob = torch.gather(rpn_score, dim=1, index=top_n_index)
|
||||
part_imgs = torch.zeros([batch, self.topN, 3, 224, 224]).cuda()
|
||||
for i in range(batch):
|
||||
for j in range(self.topN):
|
||||
[y0, x0, y1, x1] = top_n_cdds[i][j, 1:5].astype(np.int)
|
||||
part_imgs[i:i + 1, j] = F.interpolate(x_pad[i:i + 1, :, y0:y1, x0:x1], size=(224, 224), mode='bilinear',
|
||||
align_corners=True)
|
||||
part_imgs = part_imgs.view(batch * self.topN, 3, 224, 224)
|
||||
_, _, part_features = self.pretrained_model(part_imgs.detach())
|
||||
part_feature = part_features.view(batch, self.topN, -1)
|
||||
part_feature = part_feature[:, :CAT_NUM, ...].contiguous()
|
||||
part_feature = part_feature.view(batch, -1)
|
||||
# concat_logits have the shape: B*200
|
||||
concat_out = torch.cat([part_feature, feature], dim=1)
|
||||
concat_logits = self.concat_net(concat_out)
|
||||
raw_logits = resnet_out
|
||||
# part_logits have the shape: B*N*200
|
||||
part_logits = self.partcls_net(part_features).view(batch, self.topN, -1)
|
||||
return [raw_logits, concat_logits, part_logits, top_n_index, top_n_prob]
|
||||
|
||||
|
||||
def list_loss(logits, targets):
|
||||
temp = F.log_softmax(logits, -1)
|
||||
loss = [-temp[i][targets[i].item()] for i in range(logits.size(0))]
|
||||
return torch.stack(loss)
|
||||
|
||||
|
||||
def ranking_loss(score, targets, proposal_num=PROPOSAL_NUM):
|
||||
loss = Variable(torch.zeros(1).cuda())
|
||||
batch_size = score.size(0)
|
||||
for i in range(proposal_num):
|
||||
targets_p = (targets > targets[:, i].unsqueeze(1)).type(torch.cuda.FloatTensor)
|
||||
pivot = score[:, i].unsqueeze(1)
|
||||
loss_p = (1 - pivot + score) * targets_p
|
||||
loss_p = torch.sum(F.relu(loss_p))
|
||||
loss += loss_p
|
||||
return loss / batch_size
|
||||
|
||||
|
||||
# anchors.py
|
||||
_default_anchors_setting = (
|
||||
dict(layer='p3', stride=32, size=48, scale=[2 ** (1. / 3.), 2 ** (2. / 3.)], aspect_ratio=[0.667, 1, 1.5]),
|
||||
dict(layer='p4', stride=64, size=96, scale=[2 ** (1. / 3.), 2 ** (2. / 3.)], aspect_ratio=[0.667, 1, 1.5]),
|
||||
dict(layer='p5', stride=128, size=192, scale=[1, 2 ** (1. / 3.), 2 ** (2. / 3.)], aspect_ratio=[0.667, 1, 1.5]),
|
||||
)
|
||||
|
||||
|
||||
def generate_default_anchor_maps(anchors_setting=None, input_shape=INPUT_SIZE):
|
||||
"""
|
||||
generate default anchor
|
||||
:param anchors_setting: all informations of anchors
|
||||
:param input_shape: shape of input images, e.g. (h, w)
|
||||
:return: center_anchors: # anchors * 4 (oy, ox, h, w)
|
||||
edge_anchors: # anchors * 4 (y0, x0, y1, x1)
|
||||
anchor_area: # anchors * 1 (area)
|
||||
"""
|
||||
if anchors_setting is None:
|
||||
anchors_setting = _default_anchors_setting
|
||||
|
||||
center_anchors = np.zeros((0, 4), dtype=np.float32)
|
||||
edge_anchors = np.zeros((0, 4), dtype=np.float32)
|
||||
anchor_areas = np.zeros((0,), dtype=np.float32)
|
||||
input_shape = np.array(input_shape, dtype=int)
|
||||
|
||||
for anchor_info in anchors_setting:
|
||||
|
||||
stride = anchor_info['stride']
|
||||
size = anchor_info['size']
|
||||
scales = anchor_info['scale']
|
||||
aspect_ratios = anchor_info['aspect_ratio']
|
||||
|
||||
output_map_shape = np.ceil(input_shape.astype(np.float32) / stride)
|
||||
output_map_shape = output_map_shape.astype(np.int)
|
||||
output_shape = tuple(output_map_shape) + (4,)
|
||||
ostart = stride / 2.
|
||||
oy = np.arange(ostart, ostart + stride * output_shape[0], stride)
|
||||
oy = oy.reshape(output_shape[0], 1)
|
||||
ox = np.arange(ostart, ostart + stride * output_shape[1], stride)
|
||||
ox = ox.reshape(1, output_shape[1])
|
||||
center_anchor_map_template = np.zeros(output_shape, dtype=np.float32)
|
||||
center_anchor_map_template[:, :, 0] = oy
|
||||
center_anchor_map_template[:, :, 1] = ox
|
||||
for scale in scales:
|
||||
for aspect_ratio in aspect_ratios:
|
||||
center_anchor_map = center_anchor_map_template.copy()
|
||||
center_anchor_map[:, :, 2] = size * scale / float(aspect_ratio) ** 0.5
|
||||
center_anchor_map[:, :, 3] = size * scale * float(aspect_ratio) ** 0.5
|
||||
|
||||
edge_anchor_map = np.concatenate((center_anchor_map[..., :2] - center_anchor_map[..., 2:4] / 2.,
|
||||
center_anchor_map[..., :2] + center_anchor_map[..., 2:4] / 2.),
|
||||
axis=-1)
|
||||
anchor_area_map = center_anchor_map[..., 2] * center_anchor_map[..., 3]
|
||||
center_anchors = np.concatenate((center_anchors, center_anchor_map.reshape(-1, 4)))
|
||||
edge_anchors = np.concatenate((edge_anchors, edge_anchor_map.reshape(-1, 4)))
|
||||
anchor_areas = np.concatenate((anchor_areas, anchor_area_map.reshape(-1)))
|
||||
|
||||
return center_anchors, edge_anchors, anchor_areas
|
||||
|
||||
|
||||
def hard_nms(cdds, topn=10, iou_thresh=0.25):
|
||||
if not (type(cdds).__module__ == 'numpy' and len(cdds.shape) == 2 and cdds.shape[1] >= 5):
|
||||
raise TypeError('edge_box_map should be N * 5+ ndarray')
|
||||
|
||||
cdds = cdds.copy()
|
||||
indices = np.argsort(cdds[:, 0])
|
||||
cdds = cdds[indices]
|
||||
cdd_results = []
|
||||
|
||||
res = cdds
|
||||
|
||||
while res.any():
|
||||
cdd = res[-1]
|
||||
cdd_results.append(cdd)
|
||||
if len(cdd_results) == topn:
|
||||
return np.array(cdd_results)
|
||||
res = res[:-1]
|
||||
|
||||
start_max = np.maximum(res[:, 1:3], cdd[1:3])
|
||||
end_min = np.minimum(res[:, 3:5], cdd[3:5])
|
||||
lengths = end_min - start_max
|
||||
intersec_map = lengths[:, 0] * lengths[:, 1]
|
||||
intersec_map[np.logical_or(lengths[:, 0] < 0, lengths[:, 1] < 0)] = 0
|
||||
iou_map_cur = intersec_map / ((res[:, 3] - res[:, 1]) * (res[:, 4] - res[:, 2]) + (cdd[3] - cdd[1]) * (
|
||||
cdd[4] - cdd[2]) - intersec_map)
|
||||
res = res[iou_map_cur < iou_thresh]
|
||||
|
||||
return np.array(cdd_results)
|
||||
|
||||
|
||||
#### -------------------------------如何定义batch的读写方式-------------------------------
|
||||
# 默认读写方式
|
||||
def default_loader(path):
|
||||
try:
|
||||
img = Image.open(path).convert("RGB")
|
||||
if img is not None:
|
||||
return img
|
||||
except:
|
||||
print("error image:{}".format(path))
|
||||
|
||||
|
||||
def opencv_isvalid(img_path):
|
||||
img_bgr = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)
|
||||
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
||||
return img_bgr
|
||||
|
||||
|
||||
# 判断图片是否为无效
|
||||
def IsValidImage(img_path):
|
||||
vaild = True
|
||||
if img_path.endswith(".tif") or img_path.endswith(".tiff"):
|
||||
vaild = False
|
||||
return vaild
|
||||
try:
|
||||
img = opencv_isvalid(img_path)
|
||||
if img is None:
|
||||
vaild = False
|
||||
return vaild
|
||||
except:
|
||||
vaild = False
|
||||
return vaild
|
||||
|
||||
|
||||
class MyDataset(Dataset):
|
||||
def __init__(self, dir_path, transform=None, loader=default_loader):
|
||||
fh, imgs = list(), list()
|
||||
num = 0
|
||||
for root, dirs, files in os.walk(dir_path):
|
||||
for file in files:
|
||||
try:
|
||||
img_path = os.path.join(root + os.sep, file)
|
||||
num += 1
|
||||
if IsValidImage(img_path):
|
||||
fh.append(img_path)
|
||||
else:
|
||||
os.remove(img_path)
|
||||
|
||||
except:
|
||||
print("image is broken")
|
||||
print("total images is:{}".format(num))
|
||||
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
imgs.append(line)
|
||||
|
||||
self.imgs = imgs
|
||||
self.transform = transform
|
||||
self.loader = loader
|
||||
|
||||
def __getitem__(self, item):
|
||||
fh = self.imgs[item]
|
||||
img = self.loader(fh)
|
||||
if self.transform is not None:
|
||||
img = self.transform(img)
|
||||
return fh, img
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs)
|
||||
|
||||
|
||||
#### -------------------------------如何定义batch的读写方式-------------------------------
|
||||
|
||||
|
||||
#### -------------------------------图像模糊的定义-------------------------------
|
||||
def variance_of_laplacian(image):
|
||||
return cv2.Laplacian(image, cv2.CV_64f).var()
|
||||
|
||||
|
||||
## 如何定义接口函数
|
||||
def imgQualJudge(img, QA_THRESHOLD):
|
||||
'''
|
||||
:param img:
|
||||
:param QA_THRESHOLD: 越高越清晰
|
||||
:return: 是否模糊，0为模糊，1为清晰
|
||||
'''
|
||||
|
||||
norheight = 1707
|
||||
norwidth = 1280
|
||||
flag = 0
|
||||
# 筛选尺寸
|
||||
if max(img.shape[0], img.shape[1]) < 320:
|
||||
flag = '10002'
|
||||
return flag
|
||||
|
||||
# 模糊筛选部分
|
||||
if img.shape[0] <= img.shape[1]:
|
||||
size1 = (norheight, norwidth)
|
||||
timage = cv2.resize(img, size1)
|
||||
else:
|
||||
size2 = (norwidth, norheight)
|
||||
timage = cv2.resize(img, size2)
|
||||
|
||||
tgray = cv2.cvtColor(timage, cv2.COLOR_BGR2GRAY)
|
||||
halfgray = tgray[0:int(tgray.shape[0] / 2), 0:tgray.shape[1]]
|
||||
norgrayImg = np.zeros(halfgray.shape, np.int8)
|
||||
cv2.normalize(halfgray, norgrayImg, 0, 255, cv2.NORM_MINMAX)
|
||||
fm = variance_of_laplacian(norgrayImg) # 模糊值
|
||||
if fm < QA_THRESHOLD:
|
||||
flag = '10001'
|
||||
return flag
|
||||
return flag
|
||||
|
||||
|
||||
def process(img_path):
|
||||
img = Image.open(img_path).convert("RGB")
|
||||
valid = True
|
||||
low_quality = "10001"
|
||||
size_error = "10002"
|
||||
|
||||
flag = imgQualJudge(np.array(img), 5)
|
||||
if flag == low_quality or flag == size_error or not img or 0 in np.asarray(img).shape[:2]:
|
||||
valid = False
|
||||
|
||||
return valid
|
||||
|
||||
|
||||
#### -------------------------------图像模糊的定义-------------------------------
|
||||
|
||||
def build_dict():
|
||||
dict_club = dict()
|
||||
dict_club[0] = ["身份证", 0.999999]
|
||||
dict_club[1] = ["校园卡", 0.890876]
|
||||
return dict_club
|
||||
|
||||
|
||||
class Classifier():
|
||||
def __init__(self):
|
||||
self.device = torch.device('cuda')
|
||||
self.class_id_name_dict = build_dict()
|
||||
self.mean = [0.485, 0.456, 0.406]
|
||||
self.std = [0.229, 0.224, 0.225]
|
||||
self.input_size = 448
|
||||
self.use_cuda = torch.cuda.is_available()
|
||||
self.model = AttentionNet(topN=4)
|
||||
self.model.eval()
|
||||
|
||||
checkpoint = torch.load("./.ckpt")
|
||||
newweights = checkpoint['net_state_dict']
|
||||
|
||||
# 多卡测试转为单卡
|
||||
new_state_dic = OrderedDict()
|
||||
for k, v in newweights.items():
|
||||
name = k[7:] if k.startwith("module.") else k
|
||||
new_state_dic[name] = v
|
||||
|
||||
self.model.load_state_dict(new_state_dic)
|
||||
self.model = self.model.to(self.device)
|
||||
|
||||
def evalute(self, dir_path):
|
||||
data = MyDataset(dir_path, transform=self.preprocess)
|
||||
dataloader = DataLoader(dataset=data, batch_size=32, num_workers=8)
|
||||
|
||||
self.model.eval()
|
||||
with torch.no_grad():
|
||||
num = 0
|
||||
for i, (data, path) in enumerate(dataloader, 1):
|
||||
data = data.to(self.device)
|
||||
output = self.model(data)
|
||||
for j in range(len(data)):
|
||||
img_path = path[j]
|
||||
img_output = output[1][j]
|
||||
score, label, type = self.postprocess(img_output)
|
||||
out_dict, score = self.process(score, label, type)
|
||||
class_id = out_dict["results"]["class2"]["code"]
|
||||
num += 1
|
||||
if class_id != '00038':
|
||||
os.remove(img_path)
|
||||
else:
|
||||
continue
|
||||
|
||||
def preprocess(self, img):
|
||||
img = transforms.Resize((600, 600), Image.BILINEAR)(img)
|
||||
img = transforms.CenterCrop(self.input_size)(img)
|
||||
img = transforms.ToTensor()(img)
|
||||
img = transforms.Normalize(self.mean, self.std)
|
||||
|
||||
def postprocess(self, output):
|
||||
pred_logits = F.softmax(output, dim=0)
|
||||
score, label = pred_logits.max(0)
|
||||
score = score.item()
|
||||
label = label.item()
|
||||
type = self.class_id_name_dict[label][0]
|
||||
return score, label, type
|
||||
|
||||
def process(self, score, label, type):
|
||||
success_code = "200"
|
||||
lower_conf_code = "10008"
|
||||
|
||||
threshold = float(self.class_id_name_dict[label][1])
|
||||
if threshold > 0.99:
|
||||
threshold = 0.99
|
||||
if threshold < 0.9:
|
||||
threshold = 0.9
|
||||
## 设置查勘图片较低的阈值
|
||||
if label == 38:
|
||||
threshold = 0.5
|
||||
|
||||
if score > threshold:
|
||||
status_code = success_code
|
||||
pred_label = str(label).zfill(5)
|
||||
print("pred_label:", pred_label)
|
||||
return {"code:": status_code, "message": '图像分类成功',
|
||||
"results": {"class2": {'code': pred_label, 'name': type}}}, score
|
||||
else:
|
||||
status_code = lower_conf_code
|
||||
return {"code:": status_code, "message": '图像分类置信度低，不返回结果',
|
||||
"results": {"class2": {'code': '', 'name': ''}}}, score
|
||||
|
||||
|
||||
def class_results(img_dir):
|
||||
Classifier().evalute(img_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
59
utils/config.py
Normal file
59
utils/config.py
Normal file
@ -0,0 +1,59 @@
|
||||
from yacs.config import CfgNode as CfgNode
|
||||
_C = CfgNode()
|
||||
cfg = _C
|
||||
|
||||
_C.RESIZE = 648
|
||||
|
||||
#Monitoring table of commodity identification System
|
||||
_C.MONITORPATH = '../module/ieemoo-ai-search/model/now/monitor.json'
|
||||
|
||||
_C.VIDEOPATH = '../module/ieemoo-ai-search/videos'
|
||||
|
||||
#The number of results retrieved
|
||||
_C.NUM_RESULT = 5
|
||||
|
||||
#RAFT img numbers
|
||||
_C.NUM_RAFT = 10#6 #real nu: nu-1
|
||||
|
||||
#Path for storing video images
|
||||
_C.TEST_IMG_DIR = '../module/ieemoo-ai-search/imgs'
|
||||
_C.ORIIMG = '../module/ieemoo-ai-search/imgs'
|
||||
|
||||
_C.Tempimg = '../module/ieemoo-ai-search/data'#扫码商品展示
|
||||
_C.Tempvideos = '../module/ieemoo-ai-search/tempvideos'#扫码商品展示
|
||||
|
||||
_C.DIM = 2048
|
||||
_C.flag = False
|
||||
|
||||
#angle of rotation
|
||||
_C.ANGLES = [45,90,270,315] #45,90,135,180,225,270,315
|
||||
|
||||
#Weight of feature extraction
|
||||
#_C.NETWORK = '../../module/ieemoo-ai-search/model/now/model_best.pth' #retrieval_feature 测试
|
||||
_C.NETWORK = '../module/ieemoo-ai-search/model/now/model_best.pth'
|
||||
|
||||
#Weight of RAFT
|
||||
_C.RAFTMODEL= '../module/ieemoo-ai-search/model/now/raft-things.pth'
|
||||
|
||||
_C.DEVICE = 0
|
||||
|
||||
#Similarity threshold
|
||||
_C.THRESHOLD = 0.9
|
||||
|
||||
#mask img
|
||||
_C.MASKIMG = '../module/ieemoo-ai-search/model/now/masking.jpg'
|
||||
_C.MASKIMG_old = '../module/ieemoo-ai-search/model/now/masking_old.jpg'
|
||||
|
||||
#fgbg mask img
|
||||
_C.fgbgmask = '../module/ieemoo-ai-search/model/now/ori.jpg'
|
||||
_C.fgbgmask_old = '../module/ieemoo-ai-search/model/now/ori_old.jpg'
|
||||
|
||||
_C.URL = 'https://api.ieemoo.com/emoo-api/intelligence' #online
|
||||
#_C.URL = 'http://api.test.ieemoo.com/emoo-api/intelligence'
|
||||
|
||||
_C.Vre = 'http://api.test.ieemoo.com/emoo-api/intelligence/queryVideoCompareResult.do'
|
||||
#_C.Vre = 'http://192.168.1.98:8088/emoo-api/intelligence/queryVideoCompareResult.do'
|
||||
|
||||
_C.Ocrimg = '../module/ieemoo-ai-ocr/imgs'#post ocr img
|
||||
_C.Ocrtxt = '../module/ieemoo-ai-ocr/document'#post ocr txts
|
||||
_C.Ocrvideo = '../module/ieemoo-ai-ocr/videos'#post ocr video
|
||||
121
utils/forsegmentation.py
Normal file
121
utils/forsegmentation.py
Normal file
@ -0,0 +1,121 @@
|
||||
import requests
|
||||
from base64 import b64encode
|
||||
from json import dumps
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pdb
|
||||
max_nu_area= 0
|
||||
|
||||
def get_object_mask(frame, ori_mask, mask_path, all_nu, result_path):
|
||||
global max_nu_area
|
||||
maskimg = cv2.imread(mask_path, 0)
|
||||
kernel = np.ones((5, 5), np.uint8)
|
||||
dst = ori_mask
|
||||
dst = cv2.erode(dst, kernel)
|
||||
dst = cv2.dilate(dst, kernel)
|
||||
dst = cv2.medianBlur(dst,3)
|
||||
if (cv2.__version__).split('.')[0] == '3':
|
||||
_, contours, _ = cv2.findContours(dst, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
||||
else:
|
||||
contours, _ = cv2.findContours(dst, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
|
||||
area,mask_area = 0,0
|
||||
if len(contours) == 1 or len(contours) == 0:
|
||||
return None
|
||||
for contour in contours:
|
||||
area_now = cv2.contourArea(contour)
|
||||
dst = cv2.fillPoly(dst, [contour], (0))
|
||||
if area_now > area:
|
||||
area = area_now
|
||||
Matrix = contour
|
||||
max_area = area_now
|
||||
if max_area > max_nu_area:
|
||||
max_nu_area = max_area #维护最大的五个
|
||||
flag = False
|
||||
else:flag = True
|
||||
(x, y, w, h) = cv2.boundingRect(Matrix)
|
||||
dst = cv2.fillPoly(dst, [Matrix], (255))
|
||||
coordination = [x, y, x + w, y + h]
|
||||
if max_area/(w*h)<0.3:
|
||||
return None
|
||||
#print('masking', maskimg.shape)
|
||||
#print('dst', dst.shape)
|
||||
#pdb.set_trace()
|
||||
mask_dst = cv2.bitwise_and(maskimg, dst)
|
||||
if w < 350 or h <350:
|
||||
return None
|
||||
if (cv2.__version__).split('.')[0] == '3':
|
||||
_, contours, _ = cv2.findContours(mask_dst, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
||||
else:
|
||||
contours, _ = cv2.findContours(mask_dst, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
|
||||
for contour in contours:
|
||||
mask_area_now = cv2.contourArea(contour)
|
||||
if mask_area_now > mask_area:
|
||||
mask_area = mask_area_now
|
||||
proportion = float(mask_area/max_area)
|
||||
if proportion<1.05 and proportion>0.5:
|
||||
#print(coordination)
|
||||
A,B,C = mask_dst, mask_dst, mask_dst
|
||||
mask_dst = cv2.merge([A,B,C])
|
||||
img = cv2.bitwise_and(mask_dst, frame)
|
||||
img = img[coordination[1]:coordination[3],coordination[0]:coordination[2]]
|
||||
frame = frame[coordination[1]:coordination[3],coordination[0]:coordination[2]]
|
||||
#cv2.imshow('dst',img)
|
||||
#cv2.waitKey(1000)
|
||||
#print(all_nu)
|
||||
#if all_nu>4: return 'True'
|
||||
ratio = (w/h if h>w else h/w)
|
||||
if ratio<0.5:
|
||||
return None
|
||||
if all_nu<5:
|
||||
savenu = all_nu
|
||||
elif all_nu>=5 and not(flag):
|
||||
savenu = all_nu%5
|
||||
print(savenu)
|
||||
else:
|
||||
return 'True'
|
||||
cv2.imwrite('images/' + str(savenu)+'.jpg', img)
|
||||
cv2.imwrite('images/' + 'ori'+ str(savenu)+'.jpg', frame)
|
||||
#cv2.imwrite(os.sep.join([result_path, str(all_nu)+'.jpg']), img)
|
||||
return 'con'
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def get_object_location(file_test, mask_path,result_path):
|
||||
cap = cv2.VideoCapture(file_test)
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
||||
fgbg = cv2.createBackgroundSubtractorMOG2(detectShadows = False)#高斯混合模型为基础背景
|
||||
nu= 10
|
||||
working_nu = 0
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
frame = cv2.medianBlur(frame, ksize=3)
|
||||
frame_motion = frame.copy()
|
||||
fgmask = fgbg.apply(frame_motion)
|
||||
mask= cv2.threshold(fgmask, 25, 255, cv2.THRESH_BINARY)[1] # 二值化
|
||||
mask = cv2.dilate(mask, kernel, iterations=1)
|
||||
if nu<=30:
|
||||
res = get_object_mask(frame, mask, mask_path, working_nu, result_path)
|
||||
#print(res)
|
||||
if res=='True':
|
||||
break
|
||||
elif res == 'con':
|
||||
working_nu+=1
|
||||
else: continue
|
||||
else:break
|
||||
nu+=1
|
||||
|
||||
def analysis(file_test, mask_path, result_path):
|
||||
#mask_path = 'mask.jpg'
|
||||
#result_path = 'result'
|
||||
if not (os.path.exists(result_path)):
|
||||
os.mkdir(result_path)
|
||||
get_object_location(file_test, mask_path, result_path)
|
||||
if __name__ == '__main__':
|
||||
mask_path = 'mask2.jpg'
|
||||
result_path = 'images'
|
||||
file_test = "ftp/anonymous/20210908-150524_e8e24395-fc7b-42f2-a50a-068c4ac73ee9_6921168509256.mp4"
|
||||
analysis(file_test, mask_path, result_path)
|
||||
45
utils/logging.json
Normal file
45
utils/logging.json
Normal file
@ -0,0 +1,45 @@
|
||||
{
|
||||
"version":1,
|
||||
"disable_existing_loggers":false,
|
||||
"formatters":{
|
||||
"simple":{
|
||||
"format":"%(asctime)s - %(module)s - %(thread)d - %(levelname)s : %(message)s"
|
||||
}
|
||||
},
|
||||
"handlers":{
|
||||
"console":{
|
||||
"class":"logging.StreamHandler",
|
||||
"level":"DEBUG",
|
||||
"formatter":"simple",
|
||||
"stream":"ext://sys.stdout"
|
||||
},
|
||||
"info_file_handler":{
|
||||
"class":"logging.handlers.RotatingFileHandler",
|
||||
"level":"INFO",
|
||||
"formatter":"simple",
|
||||
"filename":"../log/ieemoo-ai-search-biz.log",
|
||||
"maxBytes":10485760,
|
||||
"backupCount":20,
|
||||
"encoding":"utf8"
|
||||
},
|
||||
"error_file_handler":{
|
||||
"class":"logging.handlers.RotatingFileHandler",
|
||||
"level":"ERROR",
|
||||
"formatter":"simple",
|
||||
"filename":"../log/ieemoo-ai-search-biz.log",
|
||||
"maxBytes":10485760,
|
||||
"backupCount":20,
|
||||
"encoding":"utf8"
|
||||
}
|
||||
},
|
||||
"loggers":{
|
||||
"my_module":{
|
||||
"level":"ERROR",
|
||||
"handlers":["info_file_handler"],
|
||||
"propagate":"no"}
|
||||
},
|
||||
"root":{
|
||||
"level":"INFO",
|
||||
"handlers":["console","info_file_handler","error_file_handler"]
|
||||
}
|
||||
}
|
||||
61
utils/monitor.py
Normal file
61
utils/monitor.py
Normal file
@ -0,0 +1,61 @@
|
||||
import os
|
||||
import json
|
||||
import h5py
|
||||
import numpy as np
|
||||
from utils.config import cfg
|
||||
from cirtorch.networks.imageretrievalnet import init_network, extract_vectors
|
||||
from utils.tools import createNet,ManagingFeature
|
||||
|
||||
class Moniting:
|
||||
def __init__(self, barcode = ''):
|
||||
self.barcode = barcode
|
||||
self.jsonpath = cfg.MONITORPATH
|
||||
self.MF = ManagingFeature()
|
||||
if not os.path.exists(self.jsonpath):
|
||||
jsontext = {"monitor":[]}
|
||||
jsondata = json.dumps(jsontext)
|
||||
f = open(self.jsonpath, 'w+', encoding='utf-8')
|
||||
f.write(jsondata)
|
||||
f.close()
|
||||
|
||||
def add(self):
|
||||
with open(self.jsonpath, 'r', encoding='utf-8') as add_f:
|
||||
jsonfile = json.load(add_f)
|
||||
add_f.close()
|
||||
data = list(set(jsonfile['monitor']+self.barcode))
|
||||
jsonfile['monitor'] = data
|
||||
with open(self.jsonpath, 'w') as add_f:
|
||||
json.dump(jsonfile, add_f)
|
||||
add_f.close()
|
||||
|
||||
def search(self):
|
||||
with open(self.jsonpath, 'r', encoding='utf-8') as f:
|
||||
jsonfile = json.load(f)
|
||||
f.close()
|
||||
data = set(jsonfile['monitor'])
|
||||
if self.barcode in data:
|
||||
return 'success'
|
||||
else:
|
||||
return 'nomatch'
|
||||
|
||||
def update(self, net, transform, ms):#update monitor.json
|
||||
Dict = {}
|
||||
path = []
|
||||
collectbarcode = set()
|
||||
for name in os.listdir(cfg.IMG_DIR_TOTAL):
|
||||
barcode = name.split('_')[-1].split('.')[0]
|
||||
path.append(os.sep.join([cfg.IMG_DIR_TOTAL, name]))
|
||||
collectbarcode.add(barcode)
|
||||
vecs, img_paths = extract_vectors(net, path, cfg.RESIZE, transform, ms=ms)
|
||||
data = list(vecs.detach().cpu().numpy().T)
|
||||
for code, feature in zip(img_paths, data):
|
||||
barcode = code.split('_')[-1].split('.')[0]
|
||||
feature = feature.tolist()
|
||||
self.MF.addfeature(barcode, feature)
|
||||
Moniting(list(collectbarcode)).add()
|
||||
|
||||
if __name__ == '__main__':
|
||||
barcode = '1'
|
||||
mo = Moniting(barcode)
|
||||
print(mo.search())
|
||||
mo.add(barcode)
|
||||
71
utils/retrieval_feature.py
Normal file
71
utils/retrieval_feature.py
Normal file
@ -0,0 +1,71 @@
|
||||
# coding=utf-8
|
||||
# /usr/bin/env pythpn
|
||||
|
||||
import sys
|
||||
sys.path.append('..')
|
||||
import os
|
||||
from PIL import Image
|
||||
from cirtorch.networks.imageretrievalnet import extract_vectors, extract_vectors_o
|
||||
from utils.config import cfg
|
||||
from utils.monitor import Moniting
|
||||
import cv2 as cv
|
||||
# setting up the visible GPU
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
|
||||
|
||||
class ImageProcess():
|
||||
def __init__(self, img_dir):
|
||||
self.img_dir = img_dir
|
||||
|
||||
def process(self, uuid_barcode):
|
||||
imgs = list()
|
||||
nu = 0
|
||||
for root, dirs, files in os.walk(self.img_dir):
|
||||
for file in files:
|
||||
img_path = os.path.join(root + os.sep, file)
|
||||
try:
|
||||
image = Image.open(img_path)
|
||||
if max(image.size) / min(image.size) < 5:
|
||||
if uuid_barcode == None:
|
||||
imgs.append(img_path)
|
||||
print('\r>>>> {}/{} Train done...'.format((nu + 1),
|
||||
len(os.listdir(self.img_dir))),
|
||||
end='')
|
||||
nu+=1
|
||||
else:
|
||||
if uuid_barcode in img_path:
|
||||
imgs.append(img_path)
|
||||
except:
|
||||
print("image height/width ratio is small")
|
||||
return imgs
|
||||
|
||||
class AntiFraudFeatureDataset():
|
||||
def __init__(self, uuid_barcode=None, test_img_dir=cfg.TEST_IMG_DIR):#, model='work'):
|
||||
self.uuid_barcode = uuid_barcode
|
||||
self.TestImgDir = test_img_dir
|
||||
#self.model = model
|
||||
|
||||
def extractFeature_o(self, net, image, transform, ms):
|
||||
size = cfg.RESIZE
|
||||
#image = cv.resize(image, (size, size))
|
||||
vecs = extract_vectors_o(net, image, size,transform, ms=ms)
|
||||
feature_dict = list(vecs.detach().cpu().numpy().T)
|
||||
return feature_dict
|
||||
|
||||
def extractFeature(self, net, transform, ms):
|
||||
# extract database and query vectors
|
||||
print('>> database images...')
|
||||
images = ImageProcess(self.TestImgDir).process(self.uuid_barcode)
|
||||
#print('ori', images)
|
||||
vecs, img_paths = extract_vectors( net, images, cfg.RESIZE, transform, ms=ms)
|
||||
feature_dict = list(vecs.detach().cpu().numpy().T)
|
||||
return feature_dict
|
||||
|
||||
if __name__ == '__main__':
|
||||
from utils.tools import createNet
|
||||
net, transform, ms = createNet()
|
||||
path = '../data/imgs/1.jpg'
|
||||
image = cv.imread(path)
|
||||
affd = AntiFraudFeatureDataset()
|
||||
feature = affd.extractFeature_o(net, image, transform, ms)
|
||||
print(len(feature))
|
||||
|
||||
124
utils/tools.py
Normal file
124
utils/tools.py
Normal file
@ -0,0 +1,124 @@
|
||||
#from config import cfg
|
||||
from utils.config import cfg
|
||||
from cirtorch.networks.imageretrievalnet import init_network, extract_vectors
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
import cv2
|
||||
import numpy as np
|
||||
import requests
|
||||
import os, scipy, math
|
||||
import http.client
|
||||
#http.client.HTTPConnection._http_vsn = 10
|
||||
#http.client.HTTPConnection._http_vsn_str = 'HTTP/1.0'
|
||||
|
||||
def rotate_bound(image, angle): #ratio
|
||||
(h, w) = image.shape[:2]
|
||||
(cX, cY) = (w // 2, h // 2)
|
||||
M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
|
||||
cos = np.abs(M[0, 0])
|
||||
sin = np.abs(M[0, 1])
|
||||
nW = int((h * sin) + (w * cos))
|
||||
nH = int((h * cos) + (w * sin))
|
||||
M[0, 2] += (nW / 2) - cX
|
||||
M[1, 2] += (nH / 2) - cY
|
||||
return cv2.warpAffine(image, M, (nW, nH))
|
||||
|
||||
def createNet(): #load model
|
||||
multiscale = '[1]'
|
||||
print(">> Loading network:\n>>>> '{}'".format(cfg.NETWORK))
|
||||
state = torch.load(cfg.NETWORK)
|
||||
net_params = {}
|
||||
net_params['architecture'] = state['meta']['architecture']
|
||||
net_params['pooling'] = state['meta']['pooling']
|
||||
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
|
||||
net_params['regional'] = state['meta'].get('regional', False)
|
||||
net_params['whitening'] = state['meta'].get('whitening', False)
|
||||
net_params['mean'] = state['meta']['mean']
|
||||
net_params['std'] = state['meta']['std']
|
||||
net_params['pretrained'] = False
|
||||
net = init_network(net_params)
|
||||
net.load_state_dict(state['state_dict'])
|
||||
print(">>>> loaded network: ")
|
||||
print(net.meta_repr())
|
||||
ms = list(eval(multiscale))
|
||||
print(">>>> Evaluating scales: {}".format(ms))
|
||||
if torch.cuda.is_available():
|
||||
net.cuda()
|
||||
net.eval()
|
||||
normalize = transforms.Normalize(
|
||||
mean=net.meta['mean'],
|
||||
std=net.meta['std']
|
||||
)
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
normalize
|
||||
])
|
||||
return net, transform, ms
|
||||
|
||||
class ManagingFeature: #特征增删改查
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def addfeature(self, code, feature):
|
||||
url = os.sep.join([cfg.URL, 'addImageFeatureInfo.do'])
|
||||
json = {'code':code,
|
||||
'featureVal':feature}
|
||||
r = requests.post(url=url, data=json)
|
||||
return r.text
|
||||
|
||||
def deletefeature(self, code, timeStamp): #eg: "timeStamp":"2022/02/10 17:59:59"
|
||||
url = os.sep.join([cfg.URL, 'deletImageFeatureInfo.do'])
|
||||
json = {'code':code,
|
||||
'timeStamp':timeStamp}
|
||||
r = requests.get(url=url, params=json)
|
||||
return r.json
|
||||
|
||||
def getfeature(self, code):
|
||||
try:
|
||||
url = os.sep.join([cfg.URL, 'getImageFeatureInfo.do'])
|
||||
json = {'code': code}
|
||||
r = requests.get(url=url, params=json)
|
||||
data = r.json()['data']
|
||||
return data
|
||||
except Exception as e:
|
||||
print('>>>>>get feature error<<<<<<')
|
||||
|
||||
class EvaluteMap():
|
||||
def __init__(self):
|
||||
self.MF = ManagingFeature()
|
||||
|
||||
def match_feature(self, features, search_r):
|
||||
alldict = []
|
||||
#search_r = self.MF.getfeature(barcode)
|
||||
for feature in features:
|
||||
for r in search_r:
|
||||
dist = np.linalg.norm(feature - r)
|
||||
#alldict.append(math.pow(dist, 2))
|
||||
alldict.append(dist)
|
||||
meandist = scipy.mean(sorted(alldict)[0:5])
|
||||
return meandist
|
||||
|
||||
def match_feature_single(self, feature, search_r):
|
||||
alldict = []
|
||||
for r in search_r:
|
||||
r = np.array(r)
|
||||
feature = np.array(feature)
|
||||
dist = np.linalg.norm(feature-r)
|
||||
alldict.append(dist)
|
||||
meandist = scipy.mean(sorted(alldict)[0:2])
|
||||
return meandist
|
||||
|
||||
def match_images(self, feature_dict, barcode, choose = False, mod = 'batch'):
|
||||
if mod == 'batch':
|
||||
result = self.match_feature(feature_dict, barcode)
|
||||
return result
|
||||
else:
|
||||
result = self.match_feature_single(feature_dict, barcode)
|
||||
return result
|
||||
|
||||
if __name__ == '__main__':
|
||||
pass
|
||||
# MF = ManagingFeature()
|
||||
# result = MF.getfeature('7613035075443')
|
||||
# print(result)
|
||||
|
||||
71
utils/updateObs.py
Normal file
71
utils/updateObs.py
Normal file
@ -0,0 +1,71 @@
|
||||
from obs import ObsClient
|
||||
from datetime import datetime
|
||||
#from config import cfg
|
||||
from utils.config import cfg
|
||||
import os,threading
|
||||
import time as ti
|
||||
import base64,requests,cv2,shutil
|
||||
lock = threading.Lock()
|
||||
obsClient = ObsClient(
|
||||
access_key_id='LHXJC7GIC2NNUUHHTNVI',
|
||||
secret_access_key='sVWvEItrFKWPp5DxeMvX8jLFU69iXPpzkjuMX3iM',
|
||||
server='https://obs.cn-east-3.myhuaweicloud.com'
|
||||
)
|
||||
bucketName = 'ieemoo-ai'
|
||||
|
||||
def AddObs(file_path, status):
|
||||
with lock:
|
||||
if not cfg.flag:
|
||||
addobs(file_path, status)
|
||||
else:
|
||||
if status == '02':
|
||||
addobs(file_path, status)
|
||||
else:
|
||||
objectkey = os.path.basename(file_path)
|
||||
f_dist = os.sep.join([cfg.tempvideos, status+'_'+objectkey])
|
||||
shutil.move(file_path, f_dst)
|
||||
|
||||
def addobs(file_path, status): #save videos
|
||||
ti.sleep(5)
|
||||
videoUuid = os.path.basename(file_path).split('_')[1]
|
||||
json_data = {'videoUuid': videoUuid}
|
||||
resp = requests.post(url=cfg.Vre,
|
||||
data=json_data)
|
||||
status = resp.json()
|
||||
try:
|
||||
status = status['data']
|
||||
except Exception as e:
|
||||
ti.sleep(5)
|
||||
resp = requests.post(url=cfg.Vre,data=json_data)
|
||||
status = resp.json()
|
||||
status = status['data']
|
||||
|
||||
objectkey = os.path.basename(file_path)
|
||||
time = os.path.basename(file_path).split('-')[0]
|
||||
if objectkey.split('.')[-1] in ['avi','mp4']:
|
||||
objectkey = 'videos/'+time+'/'+status+'/'+status+'_'+objectkey
|
||||
resp = obsClient.putFile(bucketName, objectkey, file_path)
|
||||
os.remove(file_path)
|
||||
|
||||
def Addimg(uuid_barcode):
|
||||
time = uuid_barcode.split('-')[0].split('_')[-1]
|
||||
objectkey = 'imgs/'+time+'/'+uuid_barcode+'.jpg'
|
||||
file_path = os.sep.join([cfg.Tempimg, '5_'+uuid_barcode+'.jpg'])
|
||||
if not os.path.exists(file_path):
|
||||
file_path = os.sep.join([cfg.Tempimg, '3_'+uuid_barcode+'.jpg'])
|
||||
if not os.path.exists(file_path):
|
||||
file_path = os.sep.join([cfg.Tempimg, 'ex_'+uuid_barcode+'.jpg'])
|
||||
resp = obsClient.putFile(bucketName, objectkey, file_path)
|
||||
|
||||
def Addimg_content(uuid_barcode, context):
|
||||
success, encoded_image = cv2.imencode(".jpg",context)
|
||||
context = encoded_image.tobytes()
|
||||
time = uuid_barcode.split('-')[0]
|
||||
objectkey = 'imgs/'+time+'/'+uuid_barcode+'.jpg'
|
||||
resp = obsClient.putContent(bucketName, objectkey, context)
|
||||
|
||||
if __name__ == '__main__':
|
||||
import cv2
|
||||
context = cv2.imread('/home/lc/project/ieemoo-ai-search/data/imgs/20230625-094651_37dd99b0-520d-457b-8615-efdb7f53b5b4_6907992825762.jpg')
|
||||
uuid_barcode = '20230625-094651_37dd99b0-520d-457b-8615-efdb7f53b5b4'
|
||||
Addimg_content(uuid_barcode, context)
|
||||
Reference in New Issue
Block a user