codingnplayinjeonju/funny

1import torch2import torchvision.transforms as transforms3 4from torchvision.models import resnet50, ResNet50_Weights5 6# 사전 학습된 ResNet 불러오기7from torchvision import models8 9model = models.resnet50(weights=ResNet50_Weights.DEFAULT)10model.eval()11 12# 라벨 불러오기 (ImageNet 클래스 라벨)13with open("imagenet_classes.txt") as f:14    labels = [line.strip() for line in f.readlines()]15 16# 변환 파이프라인 (크기 조정, 텐서 변환, 정규화 등)17transform = transforms.Compose([18    transforms.Resize((224, 224)),19    transforms.ToTensor(),20    transforms.Normalize(21        mean=[0.485, 0.456, 0.406],22        std=[0.229, 0.224, 0.225]23    )24])25 26def classify_image(image):27    # image: PIL Image28    img_t = transform(image)29    batch_t = torch.unsqueeze(img_t, 0)30    out = model(batch_t)31    _, index = torch.max(out, 1)32    return labels[index.item()]