Update mhr/predict_tools.py

mhr/predict_tools.py

@@ -44,14 +44,6 @@ class MHRVedioCuter:
         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 MHRVedioSimpleCuter(MHRVedioCuter):
-    def __init__(self):
-        super(MHRVedioSimpleCuter, self).__init__()
-
-    def iter(self, v):
-        for img in v:
-            yield self._cut_whole(img), self._cut_hole(img), self._cut_skill(img), label.format(idx/fps, idx)
-
 
 
 class MHRStoneRecognizeMgr:
@@ -70,40 +62,40 @@ class MHRStoneRecognizeMgr:
         self.skill_model = MyTrRecognizeNet(image_padding=2)
         self._vedio_cutter = vedio_cutter
 
+    def recognize_image(self, 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):
-            i+=1
-            if i % 30 == 0: 
-                print(vname, i)
-            #if len(results) == 1:
-            #    break
-            #whole
-            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:
-                continue
-            ret_whole = True
-            #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]
-            if len(results) > 0 and dump(results[-1]) == dump(result):
+            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: