"""
Script used to load a custom model and push it to HuggingFace.
Doc: https://huggingface.co/docs/transformers/custom_models#writing-a-custom-model
"""
import os
import json
import cv2

os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

from cetacean_classifier.configuration_cetacean_classifier import (
    CetaceanClassifierConfig,
)
from cetacean_classifier.modeling_cetacean_classifier import (
    CetaceanClassifierModelForImageClassification,
)

with open("original_model_config.json", "r") as file:
    config= json.load(file)

cetacean_config = CetaceanClassifierConfig(**config)
cetacean_classifier = CetaceanClassifierModelForImageClassification(cetacean_config)

cetacean_classifier.model.load_from_checkpoint("last.ckpt")

img = cv2.imread("tail.jpg")

out = cetacean_classifier(img)
print(out)