Create app.py

app.py

@@ -0,0 +1,45 @@
+import torch
+import torchvision.transforms as transforms
+from torchvision import models
+from PIL import Image
+import gradio as gr
+
+# 모델 클래스 정의 (예: EfficientNet)
+class MyModel(torch.nn.Module):
+    def __init__(self, num_classes):
+        super(MyModel, self).__init__()
+        self.model = models.efficientnet_b0(pretrained=False, num_classes=num_classes)
+
+    def forward(self, x):
+        return self.model(x)
+
+# 모델 초기화
+num_classes = 6  # 클래스 수에 맞게 설정
+model = MyModel(num_classes)
+model.load_state_dict(torch.load('resnet18_finetuned_2.pth'), strict=False)
+model.eval()
+
+# 데이터 변환 정의
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# 클래스 이름 리스트 정의
+class_names = ['disposablecup', 'envmark', 'label', 'mug', 'nonlabel', 'reusablecup']
+
+# 예측 함수 정의
+def predict(image):
+    image = transform(image).unsqueeze(0)  # 배치 차원 추가
+    with torch.no_grad():
+        outputs = model(image)
+        _, predicted = torch.max(outputs, 1)
+    # 클래스 인덱스를 클래스 이름으로 변환
+    label = class_names[predicted.item()]
+    print("Predicted label:", label)  # 예측된 레이블 출력
+    return label
+
+# Gradio 인터페이스 설정
+iface = gr.Interface(fn=predict, inputs=gr.Image(type='pil'), outputs='label', live=True)
+iface.launch()