debug

app.py

@@ -41,6 +41,31 @@ Rect = Tuple[int, int, int, int]
 FourPoint = Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int], Tuple[int, int]]
 
 
+reader = ReaderLite(gpu=True)
+overrides={"eval_cider": False, "beam": 5, "max_len_b": 64, "patch_image_size": 480,
+           "orig_patch_image_size": 224, "no_repeat_ngram_size": 0, "seed": 42}
+models, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
+    utils.split_paths('checkpoints/ocr_general_clean.pt'),
+    arg_overrides=overrides
+)
+
+# Move models to GPU
+for model in models:
+    model.eval()
+    if use_fp16:
+        model.half()
+    if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
+        model.cuda()
+    model.prepare_for_inference_(cfg)
+
+# Initialize generator
+generator = task.build_generator(models, cfg.generation)
+
+bos_item = torch.LongTensor([task.src_dict.bos()])
+eos_item = torch.LongTensor([task.src_dict.eos()])
+pad_idx = task.src_dict.pad()
+
+
 def four_point_transform(image: np.ndarray, rect: FourPoint) -> np.ndarray:
     (tl, tr, br, bl) = rect
 
@@ -81,10 +106,7 @@ def draw_boxes(image, bounds, color='red', width=4):
     return image
 
 
-def encode_text(task, text, length=None, append_bos=False, append_eos=False):
-    bos_item = torch.LongTensor([task.src_dict.bos()])
-    eos_item = torch.LongTensor([task.src_dict.eos()])
-
+def encode_text(text, length=None, append_bos=False, append_eos=False):
     s = task.tgt_dict.encode_line(
         line=task.bpe.encode(text),
         add_if_not_exist=False,
@@ -113,37 +135,11 @@ def patch_resize_transform(patch_image_size=480, is_document=False):
     return _patch_resize_transform
 
 
-reader = ReaderLite(gpu=True)
-
-overrides={"eval_cider": False, "beam": 5, "max_len_b": 64, "patch_image_size": 480,
-           "orig_patch_image_size": 224, "no_repeat_ngram_size": 0, "seed": 42}
-models, cfg, task = checkpoint_utils.load_model_ensemble_and_task(
-    utils.split_paths('checkpoints/ocr_general_clean.pt'),
-    arg_overrides=overrides
-)
-
-# Move models to GPU
-for model in models:
-    model.eval()
-    if use_fp16:
-        model.half()
-    if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
-        model.cuda()
-    model.prepare_for_inference_(cfg)
-
-# Initialize generator
-generator = task.build_generator(models, cfg.generation)
-
-bos_item = torch.LongTensor([task.src_dict.bos()])
-eos_item = torch.LongTensor([task.src_dict.eos()])
-pad_idx = task.src_dict.pad()
-
-
 # Construct input for caption task
-def construct_sample(task, image: Image, patch_image_size=480):
+def construct_sample(image: Image, patch_image_size=480):
     patch_image = patch_resize_transform(patch_image_size)(image).unsqueeze(0)
     patch_mask = torch.tensor([True])
-    src_text = encode_text(task, "图片上的文字是什么?", append_bos=True, append_eos=True).unsqueeze(0)
+    src_text = encode_text("图片上的文字是什么?", append_bos=True, append_eos=True).unsqueeze(0)
     src_length = torch.LongTensor([s.ne(pad_idx).long().sum() for s in src_text])
     sample = {
         "id":np.array(['42']),