CSATv2 / example.py

initial push

3cce567 11 months ago

732 Bytes

import torch
from datasets import load_dataset
from transformers import AutoImageProcessor, AutoModelForImageClassification

# 예시 데이터: 고양이 이미지
dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]

# 👉 CSATv2 모델로 교체
model_name = "Hyunil/CSATv2"

# Preprocessor + Model 로드
processor = AutoImageProcessor.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForImageClassification.from_pretrained(model_name, trust_remote_code=True)

# 전처리
inputs = processor(image, return_tensors="pt")

# 추론
with torch.no_grad():
    logits = model(**inputs).logits

pred = logits.argmax(-1).item()
print("Predicted label:", model.config.id2label[pred])

	import torch
	from datasets import load_dataset
	from transformers import AutoImageProcessor, AutoModelForImageClassification

	# 예시 데이터: 고양이 이미지
	dataset = load_dataset("huggingface/cats-image")
	image = dataset["test"]["image"][0]

	# 👉 CSATv2 모델로 교체
	model_name = "Hyunil/CSATv2"

	# Preprocessor + Model 로드
	processor = AutoImageProcessor.from_pretrained(model_name, trust_remote_code=True)
	model = AutoModelForImageClassification.from_pretrained(model_name, trust_remote_code=True)

	# 전처리
	inputs = processor(image, return_tensors="pt")

	# 추론
	with torch.no_grad():
	logits = model(**inputs).logits

	pred = logits.argmax(-1).item()
	print("Predicted label:", model.config.id2label[pred])