ning8429 commited on
Commit
fb0868b
·
verified ·
1 Parent(s): 3bf0eab

Update api_server.py

Browse files
Files changed (1) hide show
  1. api_server.py +12 -4
api_server.py CHANGED
@@ -70,7 +70,11 @@ def predict():
70
 
71
  # Make a prediction using YOLO
72
  results = model(image_data)
73
-
 
 
 
 
74
  saved_images = []
75
 
76
  # 儲存辨識後的圖片到指定資料夾
@@ -80,9 +84,13 @@ def predict():
80
  result.save(yolo_path)
81
  saved_images.append(yolo_path)
82
 
83
- # 保存裁剪後的圖片
84
- cropped_images = result.save_crop(YOLO_DIR) # save_crop usually returns list of cropped image paths
85
- saved_images.extend(cropped_images) # 將裁剪後的圖片加入到返回列表
 
 
 
 
86
 
87
  end_time = time.time()
88
  inference_time = end_time - start_time
 
70
 
71
  # Make a prediction using YOLO
72
  results = model(image_data)
73
+
74
+ # 檢查 YOLO 是否返回了有效的結果
75
+ if results is None or len(results) == 0:
76
+ return jsonify({'error': 'No results from YOLO model'}), 400
77
+
78
  saved_images = []
79
 
80
  # 儲存辨識後的圖片到指定資料夾
 
84
  result.save(yolo_path)
85
  saved_images.append(yolo_path)
86
 
87
+ # 保存裁剪後的圖片(僅當 save_crop 返回有效的路徑時才加入)
88
+ cropped_images = result.save_crop(YOLO_DIR) # 有些 YOLO 版本 save_crop 不返回值
89
+ if cropped_images: # 確保不會對 None 進行迭代
90
+ if isinstance(cropped_images, list): # 如果它返回一個列表
91
+ saved_images.extend(cropped_images)
92
+ else:
93
+ saved_images.append(cropped_images)
94
 
95
  end_time = time.time()
96
  inference_time = end_time - start_time