Update classification.py

classification.py

@@ -1,86 +1,86 @@
-import math
-import torch
-from timm import create_model
-from torchvision import transforms
-from PIL import Image
-
-# 기본 크롭 비율
-DEFAULT_CROP_PCT = 0.875
-
-# EfficientNet-B0 모델 설정 및 로드
-weights_path =  "./weights/resnest101e.in1k_weight_Pear_classification.pt"  # 로컬 가중치 파일 경로
-model_name = "resnest101e"
-model = create_model(model_name, pretrained=False,num_classes=9)  # 사전 학습 로드 생략
-#model.classifier = torch.nn.Linear(model.classifier.in_features, 2)  # 이진 분류로 수정
-model.load_state_dict(torch.load(weights_path))  # 로컬 가중치 로드
-model.eval()  # 평가 모드 설정
-
-# 클래스 이름 수동 지정
-# class_labels = ["Abormal Pear", "Normal Pear"]  # 데이터셋에 맞게 수정
-class_labels = ["정상", "흑성병","과피얼룩","복숭아 순나방","복숭아 삼식나방","배 깍지벌레","잎말이 나방류", "기타","과피흑변"]
-
-# 전처리 함수
-def transforms_imagenet_eval(
-    img_path: str,
-    img_size: int = 224,
-    crop_pct: float = DEFAULT_CROP_PCT,
-    mean: tuple = (0.485, 0.456, 0.406),
-    std: tuple = (0.229, 0.224, 0.225),
-    normalize: bool = True,
-):
-    """
-    ImageNet 스타일 이미지 전처리 함수.
-
-    Args:
-        img_path (str): 이미지 경로.
-        img_size (int): 크롭 크기.
-        crop_pct (float): 크롭 비율.
-        mean (tuple): 정규화 평균.
-        std (tuple): 정규화 표준편차.
-        normalize (bool): 정규화 여부.
-
-    Returns:
-        torch.Tensor: 전처리된 이미지 텐서.
-    """
-    img = Image.open(img_path).convert("RGB")  # 이미지 로드 및 RGB 변환
-    scale_size = math.floor(img_size / crop_pct)  # 리사이즈 크기 계산
-
-    # Transform 설정
-    tfl = [
-        transforms.Resize((scale_size, scale_size), interpolation=transforms.InterpolationMode.BILINEAR),
-        transforms.CenterCrop(img_size),
-        transforms.ToTensor(),
-    ]
-    if normalize:
-        tfl += [transforms.Normalize(mean=torch.tensor(mean), std=torch.tensor(std))]
-
-    transform = transforms.Compose(tfl)
-    return transform(img)
-
-# 추론 함수
-def predict(image_path: str):
-    """
-    주어진 이미지 파일 경로를 받아 모델 추론을 수행.
-
-    Args:
-        image_path (str): 입력 이미지 경로.
-
-    Returns:
-        str: 모델 예측 결과.
-    """
-    # 이미지 전처리
-    input_tensor = transforms_imagenet_eval(
-        img_path=image_path,
-        img_size=224,
-        normalize=True
-    ).unsqueeze(0)  # 배치 차원 추가
-
-    # 모델 추론
-    with torch.no_grad():
-        prediction = model(input_tensor)
-    
-    probs = torch.nn.functional.softmax(prediction[0], dim=-1)
-    confidences = {class_labels[i]: float(probs[i]) for i in range(9)}
-
-    # 예측 결과 반환
+import math
+import torch
+from timm import create_model
+from torchvision import transforms
+from PIL import Image
+
+# 기본 크롭 비율
+DEFAULT_CROP_PCT = 0.875
+
+# EfficientNet-B0 모델 설정 및 로드
+weights_path =  "./weights/resnest101e.in1k_weight_Pear_classification.pt"  # 로컬 가중치 파일 경로
+model_name = "resnest101e"
+model = create_model(model_name, pretrained=False,num_classes=9)  # 사전 학습 로드 생략
+#model.classifier = torch.nn.Linear(model.classifier.in_features, 2)  # 이진 분류로 수정
+model.load_state_dict(torch.load(weights_path, map_location=torch.device('cpu')))  # 로컬 가중치 로드
+model.eval()  # 평가 모드 설정
+
+# 클래스 이름 수동 지정
+# class_labels = ["Abormal Pear", "Normal Pear"]  # 데이터셋에 맞게 수정
+class_labels = ["정상", "흑성병","과피얼룩","복숭아 순나방","복숭아 삼식나방","배 깍지벌레","잎말이 나방류", "기타","과피흑변"]
+
+# 전처리 함수
+def transforms_imagenet_eval(
+    img_path: str,
+    img_size: int = 224,
+    crop_pct: float = DEFAULT_CROP_PCT,
+    mean: tuple = (0.485, 0.456, 0.406),
+    std: tuple = (0.229, 0.224, 0.225),
+    normalize: bool = True,
+):
+    """
+    ImageNet 스타일 이미지 전처리 함수.
+
+    Args:
+        img_path (str): 이미지 경로.
+        img_size (int): 크롭 크기.
+        crop_pct (float): 크롭 비율.
+        mean (tuple): 정규화 평균.
+        std (tuple): 정규화 표준편차.
+        normalize (bool): 정규화 여부.
+
+    Returns:
+        torch.Tensor: 전처리된 이미지 텐서.
+    """
+    img = Image.open(img_path).convert("RGB")  # 이미지 로드 및 RGB 변환
+    scale_size = math.floor(img_size / crop_pct)  # 리사이즈 크기 계산
+
+    # Transform 설정
+    tfl = [
+        transforms.Resize((scale_size, scale_size), interpolation=transforms.InterpolationMode.BILINEAR),
+        transforms.CenterCrop(img_size),
+        transforms.ToTensor(),
+    ]
+    if normalize:
+        tfl += [transforms.Normalize(mean=torch.tensor(mean), std=torch.tensor(std))]
+
+    transform = transforms.Compose(tfl)
+    return transform(img)
+
+# 추론 함수
+def predict(image_path: str):
+    """
+    주어진 이미지 파일 경로를 받아 모델 추론을 수행.
+
+    Args:
+        image_path (str): 입력 이미지 경로.
+
+    Returns:
+        str: 모델 예측 결과.
+    """
+    # 이미지 전처리
+    input_tensor = transforms_imagenet_eval(
+        img_path=image_path,
+        img_size=224,
+        normalize=True
+    ).unsqueeze(0)  # 배치 차원 추가
+
+    # 모델 추론
+    with torch.no_grad():
+        prediction = model(input_tensor)
+    
+    probs = torch.nn.functional.softmax(prediction[0], dim=-1)
+    confidences = {class_labels[i]: float(probs[i]) for i in range(9)}
+
+    # 예측 결과 반환
     return confidences