Update app.py

app.py

@@ -18,16 +18,23 @@ def detect_objects(image):
     image_tensor = image_tensor.float() / 255.0  # 将像素值归一化到 [0, 1]
     
     # 使用模型进行推理
-    outputs = torch_model(image_tensor.unsqueeze(0))  # 添加批次维度
+    with torch.no_grad():  # 禁用梯度计算以提高性能
+        outputs = torch_model(image_tensor.unsqueeze(0))  # 添加批次维度
+
+    # 获取预测结果（根据具体的输出格式进行调整）
+    predictions = outputs['logits'] if 'logits' in outputs else outputs[0]  # 确保获取正确的输出
 
     # 格式化输出结果
     detections = []
-    for output in outputs:
-        detections.append({
-            "label": output['label'],
-            "confidence": round(output['score'], 3),
-            "box": output['box'].tolist()  # 转换为列表格式以便 JSON 输出
-        })
+    for i in range(predictions.shape[1]):  # 遍历每个预测
+        score = predictions[0, i, -1].item()  # 假设最后一个维度是分数
+        if score > 0.5:  # 设置阈值过滤低置信度的预测
+            box = predictions[0, i, :-1].tolist()  # 获取边界框坐标（假设在前面）
+            detections.append({
+                "label": f"Object {i}",  # 根据实际情况替换标签
+                "confidence": round(score, 3),
+                "box": box,
+            })
     
     return detections