ds3lab / docparser Goto Github PK

import argparse
import json
import logging
import os
import cv2

import skimage.io

from docparser import stage1_entity_detector


image_name="outputname-01.png"
sample_img_path = "/home/ubuntu/projects/pdf-expirements/DocParser/images/{}".format(image_name)
output_dir = "/home/ubuntu/projects/pdf-expirements/DocParser/demos/output"

entity_detector = stage1_entity_detector.EntityDetector()
entity_detector.init_model(default_weights="highlevel_wsft")

sample_img = skimage.io.imread(sample_img_path)
results = entity_detector.predict(sample_img)

image = cv2.imread(sample_img_path)

shapes = []
for pred in results["prediction_list"]:
  box = pred["bbox_orig_coords"].tolist()
  label = pred["class_name"]
  score = pred["pred_score"]
  shapes.append({"box": box, "label": label, "score": score})

for shape in shapes:
  y1, x1, y2, x2 = shape["box"]
  image = cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)


cv2.imwrite("{}/{}".format(output_dir, image_name), image)