import json
from PIL import Image
from argparse import ArgumentParser
from ecos_core.indad.indad_onnx import Model
from ecos_core.indad.sample_code.data_transform import transform_default

if __name__ == "__main__":
    opt_path = "./opt.json"
    
    onnx_path = "./indad.onnx"
    img_path = "./709499392111057317098762511116.png"

    transform = transform_default
    img = Image.open(img_path)
    img = img.convert("RGB")
    img_trans = transform(img).unsqueeze(dim=0)
    parser = ArgumentParser()
    opts = parser.parse_args()
    with open(opt_path, 'r') as f:
        opts.__dict__ = json.load(f)
    model = Model(opts)
    #Build model with custom weight
    model.build_model(custom_weight=onnx_path)
    img_score, pxl_score = model(img_trans)
    print(img_score, pxl_score)
