Spaces:
Sleeping
Sleeping
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| import os | |
| import sys | |
| sys.path.insert( | |
| 0, | |
| os.path.abspath( | |
| os.path.join( | |
| os.path.dirname( | |
| os.path.abspath(__file__)), | |
| '../../'))) | |
| from deepdoc.vision.seeit import draw_box | |
| from deepdoc.vision import OCR, init_in_out | |
| import argparse | |
| import numpy as np | |
| def main(args): | |
| ocr = OCR() | |
| images, outputs = init_in_out(args) | |
| for i, img in enumerate(images): | |
| bxs = ocr(np.array(img)) | |
| bxs = [(line[0], line[1][0]) for line in bxs] | |
| bxs = [{ | |
| "text": t, | |
| "bbox": [b[0][0], b[0][1], b[1][0], b[-1][1]], | |
| "type": "ocr", | |
| "score": 1} for b, t in bxs if b[0][0] <= b[1][0] and b[0][1] <= b[-1][1]] | |
| img = draw_box(images[i], bxs, ["ocr"], 1.) | |
| img.save(outputs[i], quality=95) | |
| with open(outputs[i] + ".txt", "w+") as f: | |
| f.write("\n".join([o["text"] for o in bxs])) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--inputs', | |
| help="Directory where to store images or PDFs, or a file path to a single image or PDF", | |
| required=True) | |
| parser.add_argument('--output_dir', help="Directory where to store the output images. Default: './ocr_outputs'", | |
| default="./ocr_outputs") | |
| args = parser.parse_args() | |
| main(args) | |