diff --git a/inference.py b/inference.py index 11ef74d..7768ab7 100644 --- a/inference.py +++ b/inference.py @@ -122,9 +122,9 @@ def main(): if not os.path.exists(args.image_path): raise Exception("Image not found") device = torch.device('cuda:0' if args.cuda else 'cpu') - image = cv2.imread(args.image_path) + height, width, _ = image.shape - ocr_results = generate_tesseract_results.get_tesseract_results(image) + ocr_results = generate_tesseract_results.get_tesseract_results(args.image_path) vocab, class_mapping = load_saved_vocab(args.saved_path) candidates = extract_candidates.get_candidates(ocr_results) candidates_with_neighbours = attach_neighbour_candidates(width, height, ocr_results, candidates) @@ -161,4 +161,4 @@ def main(): if __name__ == '__main__': - main() \ No newline at end of file + main()