Create app.py

app.py

@@ -0,0 +1,34 @@
+# app.py
+import numpy as np
+from medsam import MedSAM
+import SimpleITK as sitk
+from flask import Flask, request, jsonify
+
+app = Flask(__name__)
+model = MedSAM.load_from_checkpoint("/app/medsam_vit_b.pth")
+
+@app.route('/segment', methods=['POST'])
+def segment():
+    # 接收前端数据
+    data = request.json
+    image = np.array(data['image'])
+    prompt = data['prompt']  # 格式: [{'type':'point','data':[x,y]}, ...]
+    
+    # 预处理图像 (HxW -> HxWx3)
+    img_3c = np.repeat(image[:,:,None], 3, axis=-1)
+    
+    # 转换提示格式
+    medsam_prompt = []
+    for p in prompt:
+        if p['type'] == 'point':
+            medsam_prompt.append({'point': p['data'], 'label': 1})
+        elif p['type'] == 'box':
+            medsam_prompt.append({'box': p['data']})
+    
+    # 模型推理
+    mask = model.predict(img_3c, medsam_prompt)
+    
+    return jsonify({"mask": mask.tolist()})
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=7860)