added small listen server to allow inferencing (todo: allow reading from base64)

This commit is contained in:
mrq 2023-08-05 16:50:53 +00:00
parent 4a799aa991
commit ba2ca9c24d

View File

@ -3,17 +3,30 @@ from pathlib import Path
from .inference import CAPTCHA
def main():
parser = argparse.ArgumentParser("CAPTCHA")
parser.add_argument("path", type=Path)
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 = parser.parse_args()
args, unknown = parser.parse_known_args()
captcha = CAPTCHA( config=args.yaml, ckpt=args.ckpt, device=args.device )
answer = captcha.inference( path=args.path, temperature=args.temp )
print("Answer:", answer)
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()