import json, os, sys, time
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "rk3588"))
EVAL = os.environ.get("EVAL_DIR", os.path.join(HERE, "eval"))
from PIL import Image
import locate3b as L
S = json.load(open(EVAL + "/eval_set.json"))
loc = L.Locator()
out = {}
t0 = time.time()
for k, s in enumerate(S):
    img = Image.open(EVAL + "/" + s["file"])
    q = "Locate all the instances that matches the following description: %s." % s["cat"]
    t = time.time()
    try:
        text, st = loc.locate(img, q, max_new=400, rep_penalty=1.1, max_boxes=100, dedupe_iou=0.85, temperature=0.7, top_p=0.9, seed=0)
    except Exception as e:
        text, st = "ERROR %r" % e, {}
    out[str(s["id"])] = {"text": text, "sec": round(time.time() - t, 2), "stats": st}
    if k % 10 == 0:
        print(k, "%.0fs" % (time.time() - t0), flush=True)
        json.dump(out, open(EVAL + "/preds_board_sampled.json", "w"), ensure_ascii=False)
json.dump(out, open(EVAL + "/preds_board_sampled.json", "w"), ensure_ascii=False)
loc.close(); print("BOARD DONE", "%.0fs" % (time.time() - t0))