Update app.py

app.py

@@ -12,19 +12,39 @@ model = YOLO('NailongKiller.yolo11n.pt').to(device)
 
 @spaces.GPU
 def predict(img):
-    # 优化图像预处理
-    img_resized = np.array(Image.fromarray(img).resize((640, 640)))
-    # 规范化像素值到 0-1 范围
-    img_tensor = torch.tensor(img_resized, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).div(255.0).to(device)
+    # 将输入图像转换为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))
+    
+    # 转换为numpy数组并进行预测
+    img_array = np.array(padded_img)
+    
+    # 进行预测
     results = model.predict(
-        img_tensor,
-        conf=0.50,  # 提高置信度阈值
-        iou=0.45,   # 调整 IOU 阈值
-        max_det=100 # 限制最大检测数量
+        img_array,
+        conf=0.50,  # 置信度阈值
+        iou=0.45,   # IOU阈值
+        max_det=100 # 最大检测数量
     )
-    return results[0].plot()
+    
+    # 获取预测结果
+    result_img = results[0].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]]
+    
+    return result_img
 
 demo = gr.Interface(
     fn=predict,