Image Classification
Transformers
Safetensors
cetaceanet
biology
biodiversity
custom_code
cetacean-classifier / test_model_before_push.py
MalloryWittwerEPFL's picture
Fix broken model
b3201aa
raw
history blame
786 Bytes
"""
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)