-
Notifications
You must be signed in to change notification settings - Fork 1
/
weapon_recognize.py
130 lines (109 loc) · 4.13 KB
/
weapon_recognize.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
# coding=UTF-8
import os
import cv2
import numpy as np
import numpy.typing as npt
import pandas as pd
from func_cv_imread import cv_imread
from func_cv_imwrite import cv_imwrite, cv_imwrite_png
from func_img_proc import black_area, img_similarity, scale_image
# LR1541-1696 HL958-996
def cut_weapon(
img: np.ndarray[int, np.dtype[np.uint8]]
) -> np.ndarray[int, np.dtype[np.uint8]]: # 裁剪
_h = img.shape[0]
_w = img.shape[1]
assert isinstance(_h, int)
assert isinstance(_w, int)
if _h == 1080 and _w == 1920: # 1080P 16:9
return img[958:996, 1541:1696]
if _h == 1600 and _w == 2560: # 2K 16:10
return scale_image(img[1437:1488, 2054:2261], 0.75)
return img[958:996, 1541:1696]
def binary_weapon(
img: np.ndarray[int, np.dtype[np.uint8]]
) -> np.ndarray[int, np.dtype[np.uint8]]: # 二值化
return cv2.threshold(img, 225, 255, cv2.THRESH_BINARY_INV)[1] # 二值化
def dilate_weapon(
img: np.ndarray[int, np.dtype[np.uint8]]
) -> np.ndarray[int, np.dtype[np.uint8]]: # 膨胀
kernel = np.ones((3, 3), np.uint8)
dilation = cv2.dilate(img, kernel, iterations=1)
return dilation
def erode_weapon(
img: np.ndarray[int, np.dtype[np.uint8]]
) -> np.ndarray[int, np.dtype[np.uint8]]: # 腐蚀
kernel = np.ones((3, 3), np.uint8)
erosion = cv2.erode(img, kernel, iterations=1)
return erosion
def weapon_img_process(img: np.ndarray[int, np.dtype[np.uint8]]):
assert img.dtype == np.uint8
if img.ndim > 2:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = binary_weapon(cut_weapon(img))
img = erode_weapon(img)
return img
def weapon_ref_generate(): # 根据截图生成标准参考图
sourcefiledir = './Ref/Weapons_Original/' # 截图文件夹，文件需以武器名称命名
destfiledir = './Ref/Weapons/'
for source in os.listdir(sourcefiledir):
img_bgr = cv_imread(sourcefiledir + source)
img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
img_cut = cut_weapon(img)
img_cut = binary_weapon(img_cut)
img_cut = erode_weapon(img_cut)
cv2.imencode('.png', img_cut)[1].tofile(
destfiledir + source[:-4] + '.png'
) # 中文路径保存
def weapon_ref_evaluate(): # 参考互相比较
filedir = './Ref/Weapons/'
maxn = len(os.listdir(filedir))
similarity = np.zeros([maxn, maxn])
black_area = np.zeros([maxn, 1])
i = 0
j = 0
for file_img1 in os.listdir(filedir):
img1 = cv_imread(filedir + file_img1)
# img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
for file_img2 in os.listdir(filedir):
img2 = cv_imread(filedir + file_img2)
# img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
sim = img_similarity(img1, img2)
# print('{} vs. {}: {}'.format(file_img1, file_img2, sim))
similarity[i, j] = sim
j += 1
black_area[i, 0] = len(np.where(img1 == 0)[0])
i += 1
j = 0
print(np.mean(similarity))
xd = pd.DataFrame(similarity)
xd.to_excel(
'./Ref/weapon_similarity.xlsx', header=os.listdir(filedir), index=os.listdir(filedir)
)
xd = pd.DataFrame(black_area.T)
xd.to_excel('./Ref/weapon_ref_blackarea.xlsx', header=os.listdir(filedir), index=False)
def weapon_recognize(img: np.ndarray[int, np.dtype[np.uint8]]):
img_cut = weapon_img_process(img)
# img = Dilate_Weapon(img)
blackarea = black_area(img_cut)
if blackarea < 900 or blackarea > 4200:
return None, 0
reffiledir = './Ref/Weapons/'
weaponlist = os.listdir(reffiledir)
i = 0
similarity = np.empty((len(weaponlist), 1))
for filename in weaponlist:
IMG_REF = cv_imread(reffiledir + filename)
sim = img_similarity(img_cut, IMG_REF)
similarity[i, 0] = sim
# if sim > 5000:
# return filename, sim
i += 1
maxsimnum = int(np.max(similarity))
maxsimname = weaponlist[np.where(similarity == np.max(similarity))[0][0]]
if maxsimnum < 5100:
return None, maxsimnum
return maxsimname[:-4], maxsimnum
if __name__ == '__main__':
weapon_ref_generate()
# Weapon_Ref_Evaluate()