33 lines
1.1 KiB
Python
Executable File
33 lines
1.1 KiB
Python
Executable File
import argparse
|
|
from pathlib import Path
|
|
from .inference import CAPTCHA
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser("CAPTCHA", allow_abbrev=False)
|
|
parser.add_argument("--listen", action='store_true')
|
|
parser.add_argument("--port", type=int, default=9090)
|
|
parser.add_argument("--yaml", type=Path, default=None)
|
|
parser.add_argument("--ckpt", type=Path, default=None)
|
|
parser.add_argument("--temp", type=float, default=1.0)
|
|
parser.add_argument("--device", default="cuda")
|
|
args, unknown = parser.parse_known_args()
|
|
|
|
captcha = CAPTCHA( config=args.yaml, ckpt=args.ckpt, device=args.device )
|
|
if args.listen:
|
|
from simple_http_server import route, server
|
|
|
|
@route("/")
|
|
def inference( path, temperature=1.0 ):
|
|
return { "answer": captcha.inference( path=Path(path), temperature=args.temp ) }
|
|
server.start(port=args.port)
|
|
else:
|
|
parser = argparse.ArgumentParser("CAPTCHA", allow_abbrev=False)
|
|
parser.add_argument("path", type=Path)
|
|
args2, unknown = parser.parse_known_args()
|
|
|
|
answer = captcha.inference( path=args2.path, temperature=args.temp )
|
|
print("Answer:", answer)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|