evazion
07b84d73fe
Do batch prediction inside the `predict` method instead of calling `predict` once for each image.
24 lines
996 B
Python
Executable File
24 lines
996 B
Python
Executable File
#!/usr/bin/env python
|
|
|
|
import json
|
|
import click
|
|
from autotagger import Autotagger
|
|
from fastai.vision.core import PILImage
|
|
|
|
@click.command(help="Automatically generate tags for an image.")
|
|
@click.option("-t", "--threshold", default=0.01, type=float, show_default=True, help="The minimum tag confidence level.")
|
|
@click.option("-n", "--limit", default=50, type=int, show_default=True, help="The maximum number of tags to return.")
|
|
@click.option("-m", "--model", default="models/model.pth", type=click.Path(exists=True), help="The model to use.")
|
|
@click.argument("file", nargs=-1, type=click.File("rb"), required=True)
|
|
def main(file, threshold, limit, model):
|
|
autotagger = Autotagger(model)
|
|
images = [PILImage.create(f) for f in file]
|
|
predictions = autotagger.predict(images, threshold=threshold, limit=limit)
|
|
|
|
for i, tags in enumerate(predictions):
|
|
data = { "filename": file[i].name, "tags": tags }
|
|
click.echo(json.dumps(data))
|
|
|
|
if __name__ == "__main__":
|
|
main()
|