1from torchvision import transforms, models
2from PIL import Image
3import torch
4
5# 모델 로드
6model = models.resnet18()
7model.fc = torch.nn.Linear(model.fc.in_features, 2)
8model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
9model.eval()
10
11# 전처리 정의
12transform = transforms.Compose([
13 transforms.Resize((224, 224)),
14 transforms.ToTensor(),
15 transforms.Normalize([0.485, 0.456, 0.406],
16 [0.229, 0.224, 0.225])
17])
18
19# 예측
20img = Image.open("example.jpg").convert("RGB")
21x = transform(img).unsqueeze(0)
22with torch.no_grad():
23 logits = model(x)
24 pred = logits.argmax(dim=1).item()
25 print("텍스처링 됨" if pred == 1 else "텍스처링 안됨")