diff --git a/tools/infer/predict_cls.py b/tools/infer/predict_cls.py index 00f0ffc1fe3387f139789490c4e6557eebb646d0..f9c4cc0f245acd2d5bd9a8576f6ff87eb9450854 100755 --- a/tools/infer/predict_cls.py +++ b/tools/infer/predict_cls.py @@ -82,7 +82,7 @@ class TextClassifier(object): cls_res = [['', 0.0]] * img_num batch_num = self.cls_batch_num - predict_time = 0 + elapse = 0 for beg_img_no in range(0, img_num, batch_num): end_img_no = min(img_num, beg_img_no + batch_num) norm_img_batch = [] @@ -107,14 +107,14 @@ class TextClassifier(object): self.predictor.run([norm_img_batch]) prob_out = self.output_tensors[0].copy_to_cpu() cls_res = self.postprocess_op(prob_out) - elapse = time.time() - starttime + elapse += time.time() - starttime for rno in range(len(cls_res)): label, score = cls_res[rno] cls_res[indices[beg_img_no + rno]] = [label, score] if '180' in label and score > self.cls_thresh: img_list[indices[beg_img_no + rno]] = cv2.rotate( img_list[indices[beg_img_no + rno]], 1) - return img_list, cls_res, predict_time + return img_list, cls_res, elapse def main(args): @@ -132,7 +132,7 @@ def main(args): valid_image_file_list.append(image_file) img_list.append(img) try: - img_list, cls_res, predict_time = text_classifier(img_list) + img_list, cls_res, elapse = text_classifier(img_list) except Exception as e: print(e) logger.info( @@ -145,7 +145,7 @@ def main(args): for ino in range(len(img_list)): print("Predicts of %s:%s" % (valid_image_file_list[ino], cls_res[ino])) print("Total predict time for %d images, cost: %.3f" % - (len(img_list), predict_time)) + (len(img_list), elapse)) if __name__ == "__main__": diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py index 54dbb03b570b605e5c05150b5fabef4d45d4337a..749a5fcdad3cdab126c83fde8ee7f1c5abc480bd 100755 --- a/tools/infer/predict_rec.py +++ b/tools/infer/predict_rec.py @@ -115,7 +115,7 @@ class TextRecognizer(object): rec_result = self.postprocess_op(preds) for rno in range(len(rec_result)): rec_res[indices[beg_img_no + rno]] = rec_result[rno] - elapse = time.time() - starttime + elapse += time.time() - starttime return rec_res, elapse