Spaces:
Sleeping
Sleeping
File size: 4,277 Bytes
a22da3a 39804bb a22da3a 87f517a 0dd6d08 87f517a a22da3a 87f517a a22da3a 4f3c7cb |
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 |
from mhr.config import *
class MHRVedioCuter:
def __init__(self, speed_ratio=1):
## config
#self.part_pos = {
# 'pos':[(628,350),(993,565)],
# 'page': [(781,575),(848,600)],
# 'hole': [(1167,197),(1250,227)],
# 'skill': [(1010,260),(1254,600)],
#}
self.part_pos = {
'hole': [1166,200,28,26],
'skill': [1014,264,240,50],
}
self.speed_ratio = speed_ratio
def iter(self, name):
vc = cv.VideoCapture(name)
fps = vc.get(cv.CAP_PROP_FPS)
print("vedio:", vc.isOpened(), fps)
label = "00:00:{:05.2f}({})"
rval = True
idx=0
while rval:
rval, img = vc.read()
idx+=1
if rval and idx%self.speed_ratio == 0:
yield self._cut_whole(img), self._cut_hole(img), self._cut_skill(img), label.format(idx/fps, idx)
vc.release()
def _cut_whole(self, img):
#pos_w, pos_h, w, h = self.part_pos['skill']
return img
def _cut_hole(self, img):
pos_w, pos_h, w, h = self.part_pos['hole']
return [ img[pos_h:pos_h+h, pos_w+w*i:pos_w+w*i+w] for i in range(3) ]
def _cut_skill(self, img):
pos_w, pos_h, w, h = self.part_pos['skill']
return [ img[pos_h+h*i:pos_h+h*i+h, pos_w:pos_w+w] for i in range(7) ]
class MHRStoneRecognizeMgr:
def __init__(self, whole_pkl, hole_pkl, vedio_cutter):
self.mapping_hole = [0,2,1,3,4]
cp = torch.load(whole_pkl)
self.whole_model = GaborFeatureNet(num_classes=2)
self.whole_model.load_state_dict(cp['model'])
if torch.cuda.is_available():
self.whole_model = self.whole_model.cuda()
#self.whole_model = torch.load(whole_pkl)
self.hole_feat_model = GaborFeatureGen(0)
if torch.cuda.is_available():
self.hole_feat_model = self.hole_feat_model.cuda()
self.hole_model = pickle.load(open(hole_pkl, 'rb'))
self.skill_model = MyTrRecognizeNet(image_padding=2)
self._vedio_cutter = vedio_cutter
def recognize_image(self, data):
print(type(data))
data_whole = tsfm_whole4cv(data[0])
data_whole = data_whole.unsqueeze(0)
if torch.cuda.is_available():
data_whole = data_whole.cuda()
ret = self.whole_model(data_whole)
if ret[0][1] - ret[0][0] < 2:
return False, []
#new hole
data_hole = torch.cat([ tsfm_hole4cv(item).unsqueeze(0) for item in data[1] ], dim=0)
if torch.cuda.is_available():
data_hole = data_hole.cuda()
output = self.hole_feat_model(data_hole)
df = pd.DataFrame(output.tolist())
ret_hole = list(self.hole_model.predict(df))
#new skill
data_skill = torch.cat([ tsfm_skill4cv(item).unsqueeze(0) for item in data[2] ], dim=0)
ret = self.skill_model(data_skill)
ret_skill = [ (x[0][0], x[1][0][-1]) for x in filter(lambda sk: sk[0][1] > 0.9 and sk[1][1] > 0.9, ret) ]
# reuslt
result = [data[3], ret_hole, ret_skill]
return True, result
if len(results) > 0 and dump(results[-1]) == dump(result):
return
def recognize(self, vname, fname=None):
def dump(rr):
return '_'.join([ str(x) for x in rr[1] ]) + "|" + '|'.join([ x[0]+":"+x[1] for x in rr[2] ])
results = []
i=0
for data in self._vedio_cutter.iter(vname):
ok, result = self.recognize_image(data)
if not ok or (len(results) > 0 and dump(results[-1]) == dump(result)):
continue
results.append(result)
if fname:
with open(fname, 'w') as f:
for result in results:
line = result[0]
line += ','
line += ','.join([ str(x) for x in result[1] ])
line += ','
line += ','.join([ x[0]+","+x[1] for x in result[2] ])
line += '\n'
f.write(line)
return results
#mgr = MHRStoneRecognizeMgr(
# whole_pkl = whole_pkl_file,
# hole_pkl = hole_pkl_file,
# vedio_cutter = MHRVedioCuter(speed_ratio),
#) |