ning8429 commited on
Commit
2b81936
·
verified ·
1 Parent(s): 973b6f6

Update api_server.py

Browse files
Files changed (1) hide show
  1. api_server.py +24 -29
api_server.py CHANGED
@@ -53,22 +53,6 @@ app = Flask(__name__)
53
  # API route for prediction(YOLO)
54
  @app.route('/predict', methods=['POST'])
55
  def predict():
56
- """
57
- Predicts the class label of an input image.
58
-
59
- Request format:
60
- {
61
- "image": [[pixel_values_gray]]
62
- }
63
-
64
- Response format:
65
- {
66
- "label": predicted_label,
67
- "pred_proba" prediction class probability
68
- "ml-latency-ms": latency_in_milliseconds
69
- (Measures time only for ML operations preprocessing with predict)
70
- }
71
- """
72
  if 'image' not in request.files:
73
  # Handle if no file is selected
74
  return 'No file selected'
@@ -83,23 +67,34 @@ def predict():
83
  except Exception as e:
84
  return jsonify({'error': str(e)}), 400
85
 
86
- # 將圖像儲存到一個緩衝區
87
- img_io = io.BytesIO()
88
- image_data.save(img_io, 'PNG')
89
- img_io.seek(0)
90
-
91
- # 返回圖像
92
- return send_file(img_io, mimetype='image/png')
93
-
94
- # # Check image shape
95
- # if image_data.size != (28, 28):
96
- # return "Invalid image shape. Expected (28, 28), take from 'demo images' folder."
 
 
 
 
 
 
 
 
 
 
 
 
97
 
98
  # # Preprocess the image
99
  # processed_image = preprocess_image(image_data)
100
 
101
- # # Make a prediction using YOLO
102
- # results = model(image_data)
103
 
104
  # # Process the YOLO output
105
  # detections = []
 
53
  # API route for prediction(YOLO)
54
  @app.route('/predict', methods=['POST'])
55
  def predict():
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56
  if 'image' not in request.files:
57
  # Handle if no file is selected
58
  return 'No file selected'
 
67
  except Exception as e:
68
  return jsonify({'error': str(e)}), 400
69
 
70
+ # Make a prediction using YOLO
71
+ results = model(image_data)
72
+
73
+ # 準備返回多張圖像
74
+ images_io = []
75
+
76
+ for i, result_img in enumerate(results.render()): # 假設 results.render() 返回的是 PIL Image 格式的圖像
77
+ img_io = io.BytesIO()
78
+ result_img.save(img_io, 'PNG') # 儲存 YOLO 處理過的圖像到緩衝區
79
+ img_io.seek(0)
80
+ images_io.append((f'image_{i}.png', img_io)) # 使用名稱區分不同的圖像
81
+ send_file(img_io, mimetype='image/png')
82
+
83
+ # 打包多張圖像為 ZIP 文件進行返回
84
+ zip_io = io.BytesIO()
85
+ with zipfile.ZipFile(zip_io, 'w') as zip_file:
86
+ for filename, image in images_io:
87
+ zip_file.writestr(filename, image.getvalue())
88
+
89
+ zip_io.seek(0)
90
+
91
+ # 返回壓縮包
92
+ return send_file(zip_io, mimetype='application/zip', as_attachment=True, download_name='predictions.zip')
93
 
94
  # # Preprocess the image
95
  # processed_image = preprocess_image(image_data)
96
 
97
+
 
98
 
99
  # # Process the YOLO output
100
  # detections = []