Update app.py

app.py

@@ -4,11 +4,12 @@ from torchvision import models
 from PIL import Image
 import gradio as gr
 
-# 모델 클래스 정의 (예: EfficientNet)
+# 모델 클래스 정의 (ResNet18)
 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)
+        self.model = models.resnet18(pretrained=False)  # ResNet18 모델 사용
+        self.model.fc = torch.nn.Linear(self.model.fc.in_features, num_classes)  # 마지막 레이어 수정
 
     def forward(self, x):
         return self.model(x)
@@ -42,4 +43,4 @@ def predict(image):
 
 # Gradio 인터페이스 설정
 iface = gr.Interface(fn=predict, inputs=gr.Image(type='pil'), outputs='label', live=True)
-iface.launch()
+iface.launch()