def main(args):
writer = SummaryWriter(log_dir=args.tb_dir + '/{}'.format(datetime.datetime.now()))
if args.confidence_threshold is not None:
# Set score_threshold for builtin models
args.opts.append('MODEL.ROI_HEADS.SCORE_THRESH_TEST')
args.opts.append(str(args.confidence_threshold))
args.opts.append('MODEL.RETINANET.SCORE_THRESH_TEST')
args.opts.append(str(args.confidence_threshold))
dataset_name = "sacrum_test"
register_sacrum_voc(dataset_name, "../dataset", "test")
cfg = setup_cfg(args)
predictor = DefaultPredictor(cfg)
dataset_dicts = DatasetCatalog.get(dataset_name)
# samples = random.sample(dataset_dicts, args.samples)
samples = dataset_dicts
for idx, (d) in enumerate(samples):
# img = cv2.imread(d["file_name"])
img = utils.read_image(d["file_name"], format=cfg.INPUT.FORMAT)
prediction = predictor(img)
visualizer = Visualizer(img[:, :, ::-1],
metadata=MetadataCatalog.get(dataset_name),
scale=args.scale)
vis = visualizer.draw_instance_predictions(prediction["instances"].to("cpu"))
writer.add_image(d["file_name"], np.transpose(vis.get_image(), axes=[2, 0, 1]), global_step=idx)
# cv2.imshow(dataset_name, vis.get_image()[:, :, ::-1])
#
# # Exit? Press ESC
# if cv2.waitKey(0) & 0xFF == 27:
# break
print("done")
return
def train(args):
cfg = config.setup_cfg(args)
trainer = Trainer(cfg)
trainer.resume_or_load(resume=args.resume)
if cfg.TEST.AUG.ENABLED:
trainer.register_hooks([
hooks.EvalHook(0,
lambda: trainer.test_with_TTA(cfg, trainer.model))
])
return trainer.train()
def eval(args):
cfg = config.setup_cfg(args)
model = Trainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
res = Trainer.test(cfg, model)
if cfg.TEST.AUG.ENABLED:
res.update(Trainer.test_with_TTA(cfg, model))
if comm.is_main_process():
verify_results(cfg, res)
return res
def main(args):
# Register sacrum dataset
register_sacrum_voc("sacrum_train", "../dataset", "train")
register_sacrum_voc("sacrum_test", "../dataset", "test")
register_sacrum_voc("sacrum_validation", "../dataset", "validation")
# Setup model configuration
cfg = setup_cfg(args)
# cfg = default_cfg(args)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
if args.eval_only:
model = Trainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume
)
res = Trainer.test(cfg, model)
return res
# Run training process
trainer = Trainer(cfg)
trainer.resume_or_load(resume=args.resume)
return trainer.train()
parser.add_argument(
"opts",
help="modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = parse_args()
# init args
cfg = get_defaults_cfg()
cfg = setup_cfg(cfg, args.cfg, args.opts)
# load data
train_data = build_loader(mode='train', cfg=cfg).get_data()
valid_data = build_loader(mode='valid', cfg=cfg).get_data()
test_data = build_loader(mode='test', cfg=cfg).get_data()
# pre process
train = build_pre_process(data=train_data, mode='train',
cfg=cfg).get_feature()
valid = build_pre_process(data=valid_data, mode='valid',
cfg=cfg).get_feature()
test = build_pre_process(data=test_data, mode='test',
cfg=cfg).get_feature()
features = [
c for c in train.columns
def main(args, event):
# ACCESS_KEY_ID = 'AKIAXNGPLG4HNVEQXTCY'
# ACCESS_SECRET_KEY = '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'
BUCKET_NAME = os.getenv("INPUT_BUCKET",
'test-bucket-2313') or 'test-pdf-bucket1'
print("BUCKET_NAME", BUCKET_NAME)
# CSV_BUCKET = os.getenv("OUTPUT_BUCKET") or "test-pdf-bucket1"
my_bucket = s3.Bucket(BUCKET_NAME)
# csv_bucket = s3.Bucket(CSV_BUCKET)
if args.confidence_threshold is not None:
# Set score_threshold for builtin models
args.opts.append('MODEL.WEIGHTS')
args.opts.append('model/model_0015999.pth')
args.opts.append('MODEL.ROI_HEADS.SCORE_THRESH_TEST')
args.opts.append(str(args.confidence_threshold))
args.opts.append('MODEL.RETINANET.SCORE_THRESH_TEST')
args.opts.append(str(args.confidence_threshold))
cfg = setup_cfg(args)
predictor = DefaultPredictor(cfg)
if "Records" in event:
records = json.loads(event['Records'][0]['body'])
all_pdfs = []
for record in records:
all_pdfs.append(
s3.ObjectSummary(BUCKET_NAME, record['s3']['object']['key']))
else:
all_bucket_objs = my_bucket.objects.all()
all_pdfs = [
pdf for pdf in all_bucket_objs if pdf.key.lower().endswith(".pdf")
]
print("all_pdfs", all_pdfs)
print(f"Total PDFs in bucket: {len(all_pdfs)}")
for i in range(len(all_pdfs)):
pdf_file = all_pdfs[i].key
pdfName = os.path.splitext(pdf_file)[0]
print(f"\n\nProcessing PDF [{i+1}/{len(all_pdfs)}]: {pdf_file}")
folders = list(my_bucket.objects.filter(Prefix='output/'))
testFlag = False
for folder in folders:
test = folder.key.split('/')
if len(test) == 3:
if test[1] == pdfName:
testFlag = True
print(f"{pdf_file} already processed!")
break
if testFlag:
continue
s3_buck = s3.Bucket(BUCKET_NAME)
folder_name = f"output/{pdfName}"
s3_buck.put_object(Bucket=BUCKET_NAME, Key=(folder_name + '/'))
pdf_bytes, pdf_doc, pdf_totalpages = get_pdf_info(all_pdfs[i])
print(f" Total Pages: {pdf_totalpages}")
# Saving the PDF file
pdf_writer = PdfFileWriter()
for pg in range(0, pdf_totalpages):
pdf_writer.addPage(pdf_doc.getPage(pg))
pdf_processing_path = '/tmp/' + pdfName + '.pdf'
with open(pdf_processing_path, "wb") as out:
pdf_writer.write(out)
for pg in range(1, pdf_totalpages + 1):
print(f" Processing page [{pg}/{pdf_totalpages}]")
pdf_page = norm_pdf_page(pdf_doc, pg)
img = pdf_page2img(pdf_bytes, pg)
image_file_path = '/tmp/test.png'
cv2.imwrite(image_file_path, img)
img = cv2.imread(image_file_path)
os.remove(image_file_path)
angle, img = deskew(img)
print(f" Page is rotated at an angle: {angle}°")
outputs = output_detectron(img, predictor, cfg)
print(f" Total Detections: {len(outputs)}")
interesting_areas = []
for x in outputs:
[x1, y1, x2, y2] = bboxes_pdf(img, pdf_page, x)
bbox_camelot = [
",".join([str(x1), str(y1),
str(x2), str(y2)])
][0] # x1,y1,x2,y2 where (x1, y1) -> left-top and (x2, y2) -> right-bottom in PDF coordinate space
interesting_areas.append(bbox_camelot)
# try:
output_camelot = camelot.read_pdf(filepath=pdf_processing_path,
pages=str(pg),
flavor="stream",
table_areas=interesting_areas)
output_camelot = [x.df for x in output_camelot]
if len(output_camelot) == 0:
print(
f" Camelot unable to extract tables information"
)
continue
for j, db in enumerate(output_camelot):
csv_buffer = StringIO()
db.to_csv(csv_buffer, index=False)
s3_resource = boto3.resource('s3')
#s3.Bucket(CSV_BUCKET).put_object(Key=(pdfName+'/'))
bytes_to_write = db.to_csv(None).encode()
# fs = s3fs.S3FileSystem(key=ACCESS_KEY_ID, secret=ACCESS_SECRET_KEY)
fs = s3fs.S3FileSystem()
# fs = s3fs.S3FileSystem()
with fs.open(
f's3://{BUCKET_NAME}/output/{pdfName}/Page {pg} - Table {j+1}.csv',
'wb') as f:
f.write(bytes_to_write)
# except Exception as e:
# print(e)
# print(f" [Traceback] Camelot unable to extract tables information")
os.remove(pdf_processing_path)