File size: 786 Bytes
b3201aa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
"""
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 configuration_cetacean_classifier import (
CetaceanClassifierConfig,
)
from 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)
|