mrq/resnet-classifier
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
mrq 2023-08-05 16:50:53 +00:00
parent 4a799aa991
commit ba2ca9c24d
1 changed files with 18 additions and 5 deletions

21
image_classifier/__main__.py
View File

@ -3,16 +3,29 @@ from pathlib import Path
from .inference import CAPTCHA from .inference import CAPTCHA
def main(): def main():
parser = argparse.ArgumentParser("CAPTCHA") parser = argparse.ArgumentParser("CAPTCHA", allow_abbrev=False)
parser.add_argument("path", type=Path) 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("--yaml", type=Path, default=None)
parser.add_argument("--ckpt", type=Path, default=None) parser.add_argument("--ckpt", type=Path, default=None)
parser.add_argument("--temp", type=float, default=1.0) parser.add_argument("--temp", type=float, default=1.0)
parser.add_argument("--device", default="cuda") 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 ) captcha = CAPTCHA( config=args.yaml, ckpt=args.ckpt, device=args.device )
answer = captcha.inference( path=args.path, temperature=args.temp ) 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) print("Answer:", answer)
if __name__ == "__main__": if __name__ == "__main__":