wehub-resource-sync
99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
from typing import List
|
|
|
|
import click
|
|
import os
|
|
from surya.input.load import load_from_folder, load_from_file
|
|
from surya.settings import settings
|
|
|
|
|
|
class CLILoader:
|
|
def __init__(self, filepath: str, cli_options: dict, highres: bool = False):
|
|
self.page_range = cli_options.get("page_range")
|
|
if self.page_range:
|
|
self.page_range = self.parse_range_str(self.page_range)
|
|
self.filepath = filepath
|
|
self.config = cli_options
|
|
self.save_images = cli_options.get("images", False)
|
|
self.debug = cli_options.get("debug", False)
|
|
self.output_dir = cli_options.get("output_dir")
|
|
|
|
# Opt in to leaving the inference server up so later commands reuse it.
|
|
if cli_options.get("keep_server"):
|
|
settings.SURYA_INFERENCE_KEEP_ALIVE = True
|
|
|
|
self.load(highres)
|
|
|
|
@staticmethod
|
|
def common_options(fn):
|
|
fn = click.argument("input_path", type=click.Path(exists=True), required=True)(
|
|
fn
|
|
)
|
|
fn = click.option(
|
|
"--output_dir",
|
|
type=click.Path(exists=False),
|
|
required=False,
|
|
default=os.path.join(settings.RESULT_DIR, "surya"),
|
|
help="Directory to save output.",
|
|
)(fn)
|
|
fn = click.option(
|
|
"--page_range",
|
|
type=str,
|
|
default=None,
|
|
help="Page range to convert, specify comma separated page numbers or ranges. Example: 0,5-10,20",
|
|
)(fn)
|
|
fn = click.option(
|
|
"--images",
|
|
is_flag=True,
|
|
help="Save images of detected bboxes.",
|
|
default=False,
|
|
)(fn)
|
|
fn = click.option(
|
|
"--debug", "-d", is_flag=True, help="Enable debug mode.", default=False
|
|
)(fn)
|
|
fn = click.option(
|
|
"--keep_server",
|
|
is_flag=True,
|
|
default=False,
|
|
help="Keep the inference server (vllm/llama.cpp) running after this command exits so later commands reuse it instead of re-spawning.",
|
|
)(fn)
|
|
return fn
|
|
|
|
def load(self, highres: bool = False):
|
|
highres_images = None
|
|
if os.path.isdir(self.filepath):
|
|
images, names = load_from_folder(self.filepath, self.page_range)
|
|
folder_name = os.path.basename(self.filepath)
|
|
if highres:
|
|
highres_images, _ = load_from_folder(
|
|
self.filepath, self.page_range, settings.IMAGE_DPI_HIGHRES
|
|
)
|
|
else:
|
|
images, names = load_from_file(self.filepath, self.page_range)
|
|
folder_name = os.path.basename(self.filepath).split(".")[0]
|
|
if highres:
|
|
highres_images, _ = load_from_file(
|
|
self.filepath, self.page_range, settings.IMAGE_DPI_HIGHRES
|
|
)
|
|
|
|
self.images = images
|
|
self.highres_images = highres_images
|
|
self.names = names
|
|
|
|
self.result_path = os.path.abspath(os.path.join(self.output_dir, folder_name))
|
|
os.makedirs(self.result_path, exist_ok=True)
|
|
|
|
@staticmethod
|
|
def parse_range_str(range_str: str) -> List[int]:
|
|
range_lst = range_str.split(",")
|
|
page_lst = []
|
|
for i in range_lst:
|
|
if "-" in i:
|
|
start, end = i.split("-")
|
|
page_lst += list(range(int(start), int(end) + 1))
|
|
else:
|
|
page_lst.append(int(i))
|
|
page_lst = sorted(
|
|
list(set(page_lst))
|
|
) # Deduplicate page numbers and sort in order
|
|
return page_lst
|