|
|
from inference import Inference |
|
|
import argparse |
|
|
import gradio as gr |
|
|
import glob |
|
|
|
|
|
def parse_option(): |
|
|
parser = argparse.ArgumentParser('MetaFG Inference script', add_help=False) |
|
|
parser.add_argument('--cfg', type=str, metavar="FILE", help='path to config file', default="configs/MetaFG_2_224.yaml") |
|
|
|
|
|
parser.add_argument('--model-path', type=str, help="path to model data", default="./ckpt_4_mf2.pth") |
|
|
parser.add_argument('--img-size', type=int, default=384, help='path to image') |
|
|
parser.add_argument('--meta-path', default="meta.txt", type=str, help='path to meta data') |
|
|
parser.add_argument('--names-path', default="names_mf2.txt", type=str, help='path to meta data') |
|
|
args = parser.parse_args() |
|
|
return args |
|
|
|
|
|
if __name__ == '__main__': |
|
|
args = parse_option() |
|
|
|
|
|
model = Inference(config_path=args.cfg, |
|
|
model_path=args.model_path, |
|
|
names_path=args.names_path) |
|
|
|
|
|
def classify(image): |
|
|
preds = model.infer(img_path=image, meta_data_path="meta.txt").squeeze() |
|
|
print(len(model.classes)) |
|
|
print(model.classes) |
|
|
confidences = {c: float(preds[i]) for i,c in enumerate(model.classes)} |
|
|
|
|
|
return confidences |
|
|
|
|
|
gr.Interface(pfn=classify, |
|
|
inputs=gr.Image(shape=(args.img_size, args.img_size), type="pil"), |
|
|
outputs=gr.Label(num_top_classes=10), |
|
|
examples=glob.glob("./example_images/*")).launch() |
|
|
|
