├── README.md ├── net.pkl ├── testimg ├── b3261155b2ff4e359f43c0b838d0b01d.jpg ├── b3861919cbf64f23bfe0aa2d32f9c4ad.jpg ├── b435127736194d3e80e669a3630f2a79.jpg ├── b567284182eb483a8d9bdadea8314e0d.jpg ├── ba0f6681824b46568d82e82ee13c5e6f.jpg ├── ba3adb7e09a049dbabc7be4ba094bd36.jpg ├── ba3b2f0aacdd43089de7e729925bb947.jpg ├── ba5e79142ce84dfaae30c72cba8309af.jpg ├── bb0f9e2351854000a18fe97411bb354f.jpg ├── bb1d23c789b14170aedfc3d3bd7bf0a1.jpg ├── bb55272d1c4b4b44a658ef04c5a5c8a7.jpg ├── bb7694dfb8014d569868afa2a20d1409.jpg ├── bb833b2598ec4b58a2cfa18149e5fe07.jpg ├── bba81793b0664ad2941c439e982d9caf.jpg └── bbb1f7347a854826a3289fc8ef10dc0d.jpg └── use.py /README.md: -------------------------------------------------------------------------------- 1 | # 易盾图标点选 2 | 3 | 自动从图片中读取小图标顺序，通过神经网络算法定位图片中全部的图标，然后进行比对确定点选的顺序及坐标。提供一张图片即可直接算出最终点选最终结果。 4 | 5 | 通过率没有细测，用几十张图片肉眼测试，感觉上要比 50% 略高一些。项目大小仅 3M。因为定位的 yolo 算法的网络被压缩到很小，便于下载，执行速度极快。 6 | 7 | 内附少量样本直接执行 use.py 脚本即可测试，该代码会直接显示标注好的图片。接口方便，代码稍加修改就能拿到自己想要的坐标信息。 8 | 9 | ``` 10 | # 脚本开发于 python3 11 | # 依赖 pytorch：（官网找安装方式，用一个比较新的版本即可）我开发使用版本为 torch-1.4.0-cp36-cp36m-win_amd64.whl 12 | # 依赖 opencv： （用这个安装命令最稳妥 pip install opencv-contrib-python==3.4.1.15 ） 13 | 需要使用sift图像算法，所以注意安装版本。 14 | 建议使用括号内的 pip 安装方式安装，最稳。觉得安装下载慢请增加豆瓣源或清华源的pypi地址 15 | 名字内含有 contrib 的opencv，你简单理解成 opencv 的增强版就行。 16 | ``` 17 | 18 | 脚本算法已做简单兼容，兼容 cpu和gpu两个版本。 -------------------------------------------------------------------------------- /net.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/net.pkl -------------------------------------------------------------------------------- /testimg/b3261155b2ff4e359f43c0b838d0b01d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/b3261155b2ff4e359f43c0b838d0b01d.jpg -------------------------------------------------------------------------------- /testimg/b3861919cbf64f23bfe0aa2d32f9c4ad.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/b3861919cbf64f23bfe0aa2d32f9c4ad.jpg -------------------------------------------------------------------------------- /testimg/b435127736194d3e80e669a3630f2a79.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/b435127736194d3e80e669a3630f2a79.jpg -------------------------------------------------------------------------------- /testimg/b567284182eb483a8d9bdadea8314e0d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/b567284182eb483a8d9bdadea8314e0d.jpg -------------------------------------------------------------------------------- /testimg/ba0f6681824b46568d82e82ee13c5e6f.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/ba0f6681824b46568d82e82ee13c5e6f.jpg -------------------------------------------------------------------------------- /testimg/ba3adb7e09a049dbabc7be4ba094bd36.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/ba3adb7e09a049dbabc7be4ba094bd36.jpg -------------------------------------------------------------------------------- /testimg/ba3b2f0aacdd43089de7e729925bb947.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/ba3b2f0aacdd43089de7e729925bb947.jpg -------------------------------------------------------------------------------- /testimg/ba5e79142ce84dfaae30c72cba8309af.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/ba5e79142ce84dfaae30c72cba8309af.jpg -------------------------------------------------------------------------------- /testimg/bb0f9e2351854000a18fe97411bb354f.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/bb0f9e2351854000a18fe97411bb354f.jpg -------------------------------------------------------------------------------- /testimg/bb1d23c789b14170aedfc3d3bd7bf0a1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/bb1d23c789b14170aedfc3d3bd7bf0a1.jpg -------------------------------------------------------------------------------- /testimg/bb55272d1c4b4b44a658ef04c5a5c8a7.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/bb55272d1c4b4b44a658ef04c5a5c8a7.jpg -------------------------------------------------------------------------------- /testimg/bb7694dfb8014d569868afa2a20d1409.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/bb7694dfb8014d569868afa2a20d1409.jpg -------------------------------------------------------------------------------- /testimg/bb833b2598ec4b58a2cfa18149e5fe07.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/bb833b2598ec4b58a2cfa18149e5fe07.jpg -------------------------------------------------------------------------------- /testimg/bba81793b0664ad2941c439e982d9caf.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/bba81793b0664ad2941c439e982d9caf.jpg -------------------------------------------------------------------------------- /testimg/bbb1f7347a854826a3289fc8ef10dc0d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cilame/yidun_icon/c8fd2d7034e320d4384735576b982b253f625f82/testimg/bbb1f7347a854826a3289fc8ef10dc0d.jpg -------------------------------------------------------------------------------- /use.py: -------------------------------------------------------------------------------- 1 | # 开发于 python3 2 | # 依赖 pytorch：（官网找安装方式）开发使用版本为 torch-1.4.0-cp36-cp36m-win_amd64.whl 3 | # 依赖 opencv： （pip install opencv-contrib-python==3.4.1.15）需要使用sift图像算法。所以注意安装版本。 4 | 5 | # 直接执行即可测试 6 | 7 | import cv2 8 | import numpy as np 9 | import torch 10 | import torch.nn as nn 11 | 12 | import os 13 | import math 14 | 15 | USE_CUDA = True if torch.cuda.is_available() else False 16 | DEVICE = 'cuda' if USE_CUDA else 'cpu' 17 | 18 | # 旋转图片的函数，目前没有被使用到 19 | def rotate_about_center(src, angle, scale=1.): 20 | w = src.shape[1] 21 | h = src.shape[0] 22 | rangle = np.deg2rad(angle) 23 | nw = (abs(np.sin(rangle)*h) + abs(np.cos(rangle)*w))*scale 24 | nh = (abs(np.cos(rangle)*h) + abs(np.sin(rangle)*w))*scale 25 | rot_mat = cv2.getRotationMatrix2D((nw*0.5, nh*0.5), angle, scale) 26 | rot_move = np.dot(rot_mat, np.array([(nw-w)*0.5, (nh-h)*0.5,0])) 27 | rot_mat[0,2] += rot_move[0] 28 | rot_mat[1,2] += rot_move[1] 29 | return cv2.warpAffine(src, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4) 30 | 31 | # 将经过 backbone 的矩阵数据转换成坐标和分类名字 32 | def parse_y_pred(ypred, anchors, class_types, islist=False, threshold=0.2, nms_threshold=0): 33 | ceillen = 5+len(class_types) 34 | sigmoid = lambda x:1/(1+math.exp(-x)) 35 | infos = [] 36 | for idx in range(len(anchors)): 37 | if USE_CUDA: 38 | a = ypred[:,:,:,4+idx*ceillen].cpu().detach().numpy() 39 | else: 40 | a = ypred[:,:,:,4+idx*ceillen].detach().numpy() 41 | for ii,i in enumerate(a[0]): 42 | for jj,j in enumerate(i): 43 | infos.append((ii,jj,idx,sigmoid(j))) 44 | infos = sorted(infos, key=lambda i:-i[3]) 45 | def get_xyxy_clz_con(info): 46 | gap = 416/ypred.shape[1] 47 | x,y,idx,con = info 48 | gp = idx*ceillen 49 | contain = torch.sigmoid(ypred[0,x,y,gp+4]) 50 | pred_xy = torch.sigmoid(ypred[0,x,y,gp+0:gp+2]) 51 | pred_wh = ypred[0,x,y,gp+2:gp+4] 52 | pred_clz = ypred[0,x,y,gp+5:gp+5+len(class_types)] 53 | if USE_CUDA: 54 | pred_xy = pred_xy.cpu().detach().numpy() 55 | pred_wh = pred_wh.cpu().detach().numpy() 56 | pred_clz = pred_clz.cpu().detach().numpy() 57 | else: 58 | pred_xy = pred_xy.detach().numpy() 59 | pred_wh = pred_wh.detach().numpy() 60 | pred_clz = pred_clz.detach().numpy() 61 | exp = math.exp 62 | cx, cy = map(float, pred_xy) 63 | rx, ry = (cx + x)*gap, (cy + y)*gap 64 | rw, rh = map(float, pred_wh) 65 | rw, rh = exp(rw)*anchors[idx][0], exp(rh)*anchors[idx][1] 66 | clz_ = list(map(float, pred_clz)) 67 | xx = rx - rw/2 68 | _x = rx + rw/2 69 | yy = ry - rh/2 70 | _y = ry + rh/2 71 | np.set_printoptions(precision=2, linewidth=200, suppress=True) 72 | if USE_CUDA: 73 | log_cons = torch.sigmoid(ypred[:,:,:,gp+4]).cpu().detach().numpy() 74 | else: 75 | log_cons = torch.sigmoid(ypred[:,:,:,gp+4]).detach().numpy() 76 | log_cons = np.transpose(log_cons, (0, 2, 1)) 77 | for key in class_types: 78 | if clz_.index(max(clz_)) == class_types[key]: 79 | clz = key 80 | break 81 | return [xx, yy, _x, _y], clz, con, log_cons 82 | def nms(infos): 83 | if not infos: return infos 84 | def iou(xyxyA,xyxyB): 85 | ax1,ay1,ax2,ay2 = xyxyA 86 | bx1,by1,bx2,by2 = xyxyB 87 | minx, miny = max(ax1,bx1), max(ay1, by1) 88 | maxx, maxy = min(ax2,bx2), min(ay2, by2) 89 | intw, inth = max(maxx-minx, 0), max(maxy-miny, 0) 90 | areaA = (ax2-ax1)*(ay2-ay1) 91 | areaB = (bx2-bx1)*(by2-by1) 92 | areaI = intw*inth 93 | return areaI/(areaA+areaB-areaI) 94 | rets = [] 95 | infos = infos[::-1] 96 | while infos: 97 | curr = infos.pop() 98 | if rets and any([iou(r[0], curr[0]) > nms_threshold for r in rets]): 99 | continue 100 | rets.append(curr) 101 | return rets 102 | if islist: 103 | v = [get_xyxy_clz_con(i) for i in infos if i[3] > threshold] 104 | if nms_threshold: 105 | return nms(v) 106 | else: 107 | return v 108 | else: 109 | return get_xyxy_clz_con(infos[0]) 110 | 111 | class Mini(nn.Module): 112 | class ConvBN(nn.Module): 113 | def __init__(self, cin, cout, kernel_size=3, stride=1, padding=None): 114 | super().__init__() 115 | padding = (kernel_size - 1) // 2 if not padding else padding 116 | self.conv = nn.Conv2d(cin, cout, kernel_size, stride, padding, bias=False) 117 | self.bn = nn.BatchNorm2d(cout, momentum=0.01) 118 | self.relu = nn.LeakyReLU(0.1, inplace=True) 119 | def forward(self, x): 120 | return self.relu(self.bn(self.conv(x))) 121 | def __init__(self, anchors, class_types, inchennel=3): 122 | super().__init__() 123 | self.oceil = len(anchors)*(5+len(class_types)) 124 | self.model = nn.Sequential( 125 | OrderedDict([ 126 | ('ConvBN_0', self.ConvBN(inchennel, 32)), 127 | ('Pool_0', nn.MaxPool2d(2, 2)), 128 | ('ConvBN_1', self.ConvBN(32, 48)), 129 | ('Pool_1', nn.MaxPool2d(2, 2)), 130 | ('ConvBN_2', self.ConvBN(48, 64)), 131 | ('Pool_2', nn.MaxPool2d(2, 2)), 132 | ('ConvBN_3', self.ConvBN(64, 80)), 133 | ('Pool_3', nn.MaxPool2d(2, 2)), 134 | ('ConvBN_4', self.ConvBN(80, 96)), 135 | ('Pool_4', nn.MaxPool2d(2, 2)), 136 | ('ConvBN_5', self.ConvBN(96, 102)), 137 | ('ConvEND', nn.Conv2d(102, self.oceil, 1)), 138 | ]) 139 | ) 140 | def forward(self, x): 141 | return self.model(x).permute(0,2,3,1) 142 | 143 | def get_clz_rect(filename, state): 144 | net = state['net'].to(DEVICE) 145 | optimizer = state['optimizer'] 146 | anchors = state['anchors'] 147 | class_types = state['class_types'] 148 | net.eval() # 重点中的重点，被坑了一整天。 149 | npimg = cv2.imread(filename) 150 | height, width = npimg.shape[:2] 151 | npimg = cv2.cvtColor(npimg, cv2.COLOR_BGR2RGB) # [y,x,c] 152 | npimg = cv2.resize(npimg, (416, 416)) 153 | npimg_ = np.transpose(npimg, (2,1,0)) # [c,x,y] 154 | y_pred = net(torch.FloatTensor(npimg_).unsqueeze(0).to(DEVICE)) 155 | img = cv2.imread(filename) 156 | v = parse_y_pred(y_pred, anchors, class_types, islist=True, threshold=0.2, nms_threshold=0.4) 157 | ret = [] 158 | for i in v: 159 | rect, clz, con, log_cons = i 160 | rw, rh = width/416, height/416 161 | rect[0],rect[2] = int(rect[0]*rw),int(rect[2]*rw) 162 | rect[1],rect[3] = int(rect[1]*rh),int(rect[3]*rh) 163 | ret.append([clz, rect]) 164 | return ret 165 | 166 | def get_cut_img(npimg, rects): 167 | ret = [] 168 | for clz, (x1,y1,x2,y2) in rects: 169 | ret.append([clz, npimg[y1:y2,x1:x2,:], (x1,y1,x2,y2)]) 170 | return ret 171 | 172 | 173 | # 处理顺序的问题 174 | def get_flags_rects(file, state): 175 | s = cv2.imread(file) 176 | a = s[160:,0*28:1*28-6,:] 177 | b = s[160:,1*28:2*28-6,:] 178 | c = s[160:,2*28:3*28-6,:] 179 | a1, a2 = a[40:60], a[0:20] 180 | b1, b2 = b[40:60], b[0:20] 181 | c1, c2 = c[40:60], c[0:20] 182 | def get_match_lens(i1, i2): 183 | i1 = cv2.resize(i1, (int(i1.shape[1]*8), int(i1.shape[0]*8))) 184 | i2 = cv2.resize(i2, (i2.shape[1]*4, i2.shape[0]*4)) 185 | s = cv2.xfeatures2d.SIFT_create() 186 | kp1,des1 = s.detectAndCompute(i1,None) 187 | kp2,des2 = s.detectAndCompute(i2,None) 188 | bf = cv2.BFMatcher() 189 | matches = bf.knnMatch(des1, des2, k=2) 190 | good = [] 191 | DIS = .88 192 | for m,n in matches: 193 | if m.distance <= DIS * n.distance: 194 | good.append([m]) 195 | i3 = cv2.drawMatchesKnn(i1,kp1,i2,kp2,good,None) 196 | # cv2.imshow('nier', i3) 197 | # cv2.waitKey(0) 198 | return len(good) 199 | def get_flag_rect(k12, cut_imgs, st): 200 | k1, k2 = k12 201 | r = [] 202 | for clz, npimg, rect in cut_imgs: 203 | if clz == '1': 204 | r1 = get_match_lens(k1, npimg) 205 | r.append([r1, rect, st]) 206 | if clz == '2': 207 | r2 = get_match_lens(k2, npimg) 208 | r.append([r2, rect, st]) 209 | return sorted(r, key=lambda i:i[0]) 210 | v = get_cut_img(s, get_clz_rect(file, state)) 211 | rs1 = get_flag_rect([a1, a2], v, 1) 212 | rs2 = get_flag_rect([b1, b2], v, 2) 213 | rs3 = get_flag_rect([c1, c2], v, 3) 214 | rs = rs1 + rs2 + rs3 215 | r = [] 216 | t = [] 217 | v = max([j for j in rs if j[2] not in t], key=lambda i:i[0]) 218 | r.append(v) 219 | t.append(v[2]) 220 | q = [] 221 | for i in rs: 222 | if i[1] == v[1]: 223 | q.append(i) 224 | for i in q: 225 | rs.remove(i) 226 | v = max([j for j in rs if j[2] not in t], key=lambda i:i[0]) 227 | r.append(v) 228 | t.append(v[2]) 229 | q = [] 230 | for i in rs: 231 | if i[1] == v[1]: 232 | q.append(i) 233 | for i in q: 234 | rs.remove(i) 235 | v = max([j for j in rs if j[2] not in t], key=lambda i:i[0]) 236 | r.append(v) 237 | t.append(v[2]) 238 | r1, r2, r3 = sorted(r,key=lambda i:i[2]) 239 | return r1[1], r2[1], r3[1] 240 | 241 | def draw_rects(filename, rects): 242 | def drawrect(img, rect, text): 243 | cv2.rectangle(img, tuple(rect[:2]), tuple(rect[2:]), (10,250,10), 2, 1) 244 | x, y = rect[:2] 245 | cv2.putText(img, text, (x,y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (10,10,250), 1) 246 | return img 247 | img = cv2.imread(filename) 248 | for idx, rect in enumerate(rects, 1): 249 | img = drawrect(img, rect, '{}'.format(idx)) 250 | cv2.imshow('test', img) 251 | cv2.waitKey(0) 252 | cv2.destroyAllWindows() 253 | 254 | xmlpath = './testimg' 255 | v = [os.path.join(xmlpath, i) for i in os.listdir(xmlpath) if i.endswith('.jpg')] 256 | v = v[::-1] 257 | 258 | print('loading net') 259 | state = torch.load('net.pkl', map_location=torch.device(DEVICE)) 260 | # state = torch.load('net.pkl') # 如果上面的加载方式不行就去掉 map_location 参数再试试 261 | print('loading net ok.') 262 | for file in v: 263 | print(file) 264 | rects = get_flags_rects(file, state) 265 | draw_rects(file, rects) 266 | --------------------------------------------------------------------------------