Migrate to ZeroGPU and update requirements for compatibility

app.py

@@ -5,7 +5,7 @@ from PIL import Image
 import numpy as np
 import io
 import torch
-import spaces  # 导入 spaces 模块
+import spaces
 
 # 初始化 FastAPI 和模型
 app = FastAPI()
@@ -14,10 +14,10 @@ app = FastAPI()
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model = YOLO('NailongKiller.yolo11n.pt').to(device)
 
-@spaces.GPU  # 使用装饰器标记需要 GPU 的函数
+@spaces.GPU
 def predict(img):
-    # 将 numpy 数组转换为 PyTorch 张量
-    img_tensor = torch.from_numpy(img).to(device)
+    img_resized = np.array(Image.fromarray(img).resize((640, 640)))
+    img_tensor = torch.from_numpy(img_resized).permute(2, 0, 1).unsqueeze(0).to(device)
     results = model.predict(img_tensor)
     return results[0].plot()
 
@@ -38,30 +38,6 @@ demo = gr.Interface(
     cache_examples=True
 )
 
-# API 端点
-@app.post("/detect/")
-async def detect_api(file: UploadFile = File(...)):
-    contents = await file.read()
-    image = Image.open(io.BytesIO(contents))
-    image_np = np.array(image)
-    
-    # 将图像移动到 GPU
-    image_tensor = torch.from_numpy(image_np).to(device)
-    
-    results = model.predict(image_tensor)
-    result = results[0]
-    
-    detections = []
-    for box in result.boxes:
-        detection = {
-            "bbox": box.xyxy[0].tolist(),
-            "confidence": float(box.conf[0]),
-            "class": int(box.cls[0])
-        }
-        detections.append(detection)
-    
-    return {"detections": detections}
-
 # 挂载 Gradio 到 FastAPI
 app = gr.mount_gradio_app(app, demo, path="/")