Update app.py

app.py

@@ -4,44 +4,22 @@ import torch
 from torchvision import transforms
 from PIL import Image
 import torch.nn.functional as F
+from ultralytics import YOLO
 
-# 加载模型
-class SimpleModel(torch.nn.Module):
-    def __init__(self, num_classes=3):
-        super(SimpleModel, self).__init__()
-        self.fc = torch.nn.Linear(3 * 224 * 224, num_classes)
-
-    def forward(self, x):
-        x = x.view(x.size(0), -1)
-        return self.fc(x)
-
-# 初始化模型并加载权重
-model = SimpleModel(num_classes=3)
-model_path = "model.pth"
-model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
-model.eval()
-
-
-def preprocess(image: Image.Image):
-    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]),
-    ])
-    return transform(image).unsqueeze(0)
 
+# 加载模型
+model = YOLO("nabird_det_ep3.pt")
 
-def classify_image(image: Image.Image):
-    input_tensor = preprocess(image)
-    with torch.no_grad():
-        output = model(input_tensor)
-        probabilities = F.softmax(output[0], dim=0)
-    top_class = probabilities.argmax().item()
-    return f"Predicted: {class_names[top_class]} (Confidence: {probabilities[top_class]:.2f})"
+def yolo_det(image: Image.Image):
+    
+    output = model.predict(image, save=False)
+    print(output)
+    
+    return output
 
 
 app = gr.Interface(
-    fn=classify_image,  # 推理函数
+    fn=yolo_det,  # 推理函数
     inputs=gr.Image(type="pil"),  # 接收输入图片，返回 PIL 格式
     outputs="text",  # 输出分类结果
     title="Image Classification with PyTorch",