Update app.py

app.py

@@ -8,25 +8,22 @@ import numpy as np
 
 app = FastAPI()
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = YOLO('nailong_yolo11.pt').to(device)
+model = YOLO('nailong_yolo11.onnx').to(device)
 
 @spaces.GPU
 def predict(img):
     # 将输入图像转换为PIL Image对象
     input_image = Image.fromarray(img)
-    original_size = input_image.size
     
-    # 计算填充尺寸
-    max_size = max(original_size)
-    pad_w = max_size - original_size[0]
-    pad_h = max_size - original_size[1]
-    
-    # 创建方形画布并保持宽高比
-    padded_img = Image.new('RGB', (max_size, max_size), (114, 114, 114))
-    padded_img.paste(input_image, (pad_w//2, pad_h//2))
+    # 保持长宽比的情况下调整尺寸
+    w, h = input_image.size
+    scale = min(640/w, 640/h)
+    new_w, new_h = int(w * scale), int(h * scale)
+    if scale != 1:
+        input_image = input_image.resize((new_w, new_h), Image.LANCZOS)
     
     # 转换为numpy数组并进行预测
-    img_array = np.array(padded_img)
+    img_array = np.array(input_image)
     
     # 进行预测
     results = model.predict(img_array)
@@ -35,11 +32,6 @@ def predict(img):
     # 获取预测结果
     result_img = result.plot()
     
-    # 裁剪回原始尺寸
-    if pad_w > 0 or pad_h > 0:
-        result_img = result_img[pad_h//2:pad_h//2 + original_size[1], 
-                              pad_w//2:pad_w//2 + original_size[0]]
-    
     # 处理检测信息
     info = {
         "detected": len(result.boxes) > 0,
@@ -70,6 +62,12 @@ def predict(img):
         for idx, det in enumerate(info["detections"], 1):
             output_text += f"\n  目标 {idx}: {det['class']} (置信度: {det['confidence']})"
     
+    # 如果需要将结果图像缩放回原始尺寸
+    if scale != 1:
+        result_img = Image.fromarray(result_img)
+        result_img = result_img.resize((w, h), Image.LANCZOS)
+        result_img = np.array(result_img)
+    
     return result_img, output_text
 
 demo = gr.Interface(