Spaces:
Sleeping
Sleeping
| 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), | |
| #) |