img_score_and_insert2.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376
  1. import asyncio
  2. import aiohttp
  3. import aiofiles
  4. import json
  5. import os
  6. from typing import Dict, Any, Tuple, Optional, List
  7. from datetime import datetime
  8. # --- 配置区域 ---
  9. # INFERENCE_SERVICE_URL = "http://127.0.0.1:7754"
  10. # STORAGE_SERVICE_URL = "http://127.0.0.1:7755"
  11. INFERENCE_SERVICE_URL = "http://192.168.77.249:7754"
  12. STORAGE_SERVICE_URL = "http://192.168.77.249:7755"
  13. # 映射关系:后端 score_type 和 image_type 现在是一致的
  14. # 灰度图不需要映射,直接透传
  15. MAIN_IMAGE_TYPES = [
  16. "front_ring",
  17. "front_coaxial",
  18. "back_ring",
  19. "back_coaxial"
  20. ]
  21. GRAY_IMAGE_TYPES = [
  22. "front_gray",
  23. "back_gray"
  24. ]
  25. # --- 基础工具函数 ---
  26. async def call_api_with_file(
  27. session: aiohttp.ClientSession,
  28. url: str,
  29. file_path: str,
  30. params: Dict[str, Any] = None,
  31. form_fields: Dict[str, Any] = None,
  32. file_field_name: str = 'file'
  33. ) -> Tuple[int, bytes]:
  34. """通用的文件上传API调用函数"""
  35. form_data = aiohttp.FormData()
  36. if form_fields:
  37. for key, value in form_fields.items():
  38. form_data.add_field(key, str(value))
  39. # 异步读取文件
  40. if not os.path.exists(file_path):
  41. return 404, b"File not found"
  42. async with aiofiles.open(file_path, 'rb') as f:
  43. content = await f.read()
  44. form_data.add_field(
  45. file_field_name,
  46. content,
  47. filename=os.path.basename(file_path),
  48. content_type='image/jpeg'
  49. )
  50. try:
  51. async with session.post(url, data=form_data, params=params) as response:
  52. response_content = await response.read()
  53. if not response.ok:
  54. print(
  55. f" [API Error] {url} -> Status: {response.status}, Msg: {response_content.decode('utf-8')[:100]}")
  56. return response.status, response_content
  57. except Exception as e:
  58. print(f" [Conn Error] {url} -> {e}")
  59. return 503, str(e).encode()
  60. # --- 业务逻辑函数 ---
  61. async def process_main_image(
  62. session: aiohttp.ClientSession,
  63. image_path: str,
  64. score_type: str,
  65. is_reflect_card: str
  66. ) -> Dict[str, Any]:
  67. """
  68. 处理【主图片】 (Ring/Coaxial):
  69. 1. 调用转正接口
  70. 2. 调用分数推理接口
  71. """
  72. print(f" [处理中] 主图: {score_type} -> {os.path.basename(image_path)}")
  73. # 1. 获取转正后的图片
  74. rectify_url = f"{INFERENCE_SERVICE_URL}/api/card_inference/card_rectify_and_center"
  75. rectify_status, rectified_image_bytes = await call_api_with_file(
  76. session, url=rectify_url, file_path=image_path
  77. )
  78. if rectify_status >= 300:
  79. raise Exception(f"转正失败: {score_type}")
  80. # 2. 获取分数JSON
  81. score_url = f"{INFERENCE_SERVICE_URL}/api/card_score/score_inference"
  82. score_params = {
  83. "score_type": score_type,
  84. "is_reflect_card": is_reflect_card
  85. }
  86. score_status, score_json_bytes = await call_api_with_file(
  87. session,
  88. url=score_url,
  89. file_path=image_path,
  90. params=score_params
  91. )
  92. if score_status >= 300:
  93. raise Exception(f"推理分数失败: {score_type}")
  94. score_json = json.loads(score_json_bytes)
  95. return {
  96. "type": "main",
  97. "image_type": score_type, # 现在 score_type 等于 image_type
  98. "rectified_image": rectified_image_bytes,
  99. "score_json": score_json
  100. }
  101. async def create_card_record(session: aiohttp.ClientSession, card_name: str) -> int:
  102. """创建新的卡牌记录"""
  103. url = f"{STORAGE_SERVICE_URL}/api/cards/created"
  104. params = {'card_name': card_name}
  105. print(f"\n[Step 2] 创建卡组记录: '{card_name}'")
  106. async with session.post(url, params=params) as response:
  107. if response.status == 201:
  108. data = await response.json()
  109. card_id = data.get('id')
  110. print(f" -> 卡组创建成功 ID: {card_id}")
  111. return card_id
  112. else:
  113. text = await response.text()
  114. raise Exception(f"创建卡组失败: {response.status} - {text}")
  115. async def upload_main_image(
  116. session: aiohttp.ClientSession,
  117. card_id: int,
  118. processed_data: Dict[str, Any]
  119. ):
  120. """上传【主图片】的处理结果 (Rectified Image + JSON)"""
  121. image_type = processed_data['image_type']
  122. url = f"{STORAGE_SERVICE_URL}/api/images/insert/{card_id}"
  123. form_data = aiohttp.FormData()
  124. form_data.add_field('image_type', image_type)
  125. form_data.add_field('json_data_str', json.dumps(processed_data['score_json'], ensure_ascii=False))
  126. form_data.add_field(
  127. 'image',
  128. processed_data['rectified_image'],
  129. filename=f'{image_type}_rectified.jpg',
  130. content_type='image/jpeg'
  131. )
  132. async with session.post(url, data=form_data) as response:
  133. if response.status == 201:
  134. print(f" -> [主图上传成功] {image_type}")
  135. else:
  136. print(f" -> [主图上传失败] {image_type} code={response.status}")
  137. async def upload_gray_image(
  138. session: aiohttp.ClientSession,
  139. card_id: int,
  140. image_type: str,
  141. file_path: str
  142. ):
  143. """上传【灰度图片】 (Raw Image only)"""
  144. print(f" [上传中] 灰度图: {image_type} -> {os.path.basename(file_path)}")
  145. url = f"{STORAGE_SERVICE_URL}/api/images/insert/gray/{card_id}"
  146. # 灰度图接口只需要 image_type 和 file
  147. form_fields = {'image_type': image_type}
  148. status, _ = await call_api_with_file(
  149. session,
  150. url=url,
  151. file_path=file_path,
  152. form_fields=form_fields,
  153. file_field_name='image' # 注意接口接收的字段名是 image
  154. )
  155. if status == 201:
  156. print(f" -> [灰度图上传成功] {image_type}")
  157. else:
  158. print(f" -> [灰度图上传失败] {image_type} code={status}")
  159. # --- 核心控制流程 ---
  160. async def process_card_images(
  161. is_reflect: bool,
  162. # 四张主图 (Optional)
  163. path_front_ring: Optional[str] = None,
  164. path_front_coaxial: Optional[str] = None,
  165. path_back_ring: Optional[str] = None,
  166. path_back_coaxial: Optional[str] = None,
  167. # 两张灰度图 (Optional)
  168. path_front_gray: Optional[str] = None,
  169. path_back_gray: Optional[str] = None,
  170. # 模式控制
  171. strict_mode: bool = True,
  172. img_index: int = None
  173. ) -> int:
  174. """
  175. 核心处理函数
  176. :param strict_mode: 如果为 True, 则必须提供所有4张主图,否则报错。
  177. 如果为 False, 允许主图缺失。灰度图始终是可选的。
  178. """
  179. # 1. 组装输入字典,过滤 None
  180. main_inputs = {
  181. "front_ring": path_front_ring,
  182. "front_coaxial": path_front_coaxial,
  183. "back_ring": path_back_ring,
  184. "back_coaxial": path_back_coaxial
  185. }
  186. gray_inputs = {
  187. "front_gray": path_front_gray,
  188. "back_gray": path_back_gray
  189. }
  190. # 2. 严格模式检查
  191. provided_main_count = sum(1 for p in main_inputs.values() if p is not None)
  192. if strict_mode:
  193. if provided_main_count != 4:
  194. print(f"\n[错误] 严格模式开启,必须提供所有4张主图。当前提供了 {provided_main_count} 张。")
  195. return -1
  196. else:
  197. if provided_main_count == 0 and not any(gray_inputs.values()):
  198. print(f"\n[错误] 未提供任何图片,无法创建。")
  199. return -1
  200. # 3. 文件存在性检查
  201. all_paths = [p for p in main_inputs.values() if p] + [p for p in gray_inputs.values() if p]
  202. for p in all_paths:
  203. if not os.path.exists(p):
  204. print(f"[错误] 文件不存在: {p}")
  205. return -1
  206. if img_index is not None:
  207. card_name = f"{img_index} 外框修正测试 {datetime.now().strftime('%m%d_%H%M%S')}"
  208. else:
  209. card_name = f"测试图片 {datetime.now().strftime('%m%d_%H%M%S')}"
  210. is_reflect_str = "true" if is_reflect else "false"
  211. async with aiohttp.ClientSession() as session:
  212. try:
  213. # --- Step 1: 主图并发处理 (推理+转正) ---
  214. # 我们先处理主图,获得结果后再创建卡片,这样如果推理全挂了就不创建空卡片了
  215. print(f"--- 开始任务: {card_name} (Strict: {strict_mode}) ---")
  216. processing_tasks = []
  217. for img_type, path in main_inputs.items():
  218. if path:
  219. task = process_main_image(session, path, img_type, is_reflect_str)
  220. processing_tasks.append(task)
  221. processed_results = []
  222. if processing_tasks:
  223. print("[Step 1] 正在推理主图片...")
  224. # return_exceptions=False 表示如果有错直接抛出,中断流程
  225. processed_results = await asyncio.gather(*processing_tasks)
  226. # --- Step 2: 创建卡片 ---
  227. card_id = await create_card_record(session, card_name)
  228. # --- Step 3: 并发上传 (主图结果 + 灰度图文件) ---
  229. print(f"\n[Step 3] 正在上传所有数据...")
  230. upload_tasks = []
  231. # 3.1 添加主图上传任务
  232. for res in processed_results:
  233. task = upload_main_image(session, card_id, res)
  234. upload_tasks.append(task)
  235. # 3.2 添加灰度图上传任务 (直接读取文件上传)
  236. for img_type, path in gray_inputs.items():
  237. if path:
  238. task = upload_gray_image(session, card_id, img_type, path)
  239. upload_tasks.append(task)
  240. if upload_tasks:
  241. await asyncio.gather(*upload_tasks)
  242. print(f"--- 流程结束, Card ID: {card_id} ---\n")
  243. return card_id
  244. except Exception as e:
  245. print(f"\n[流程终止] 发生异常: {e}")
  246. return -1
  247. # --- 同步调用封装 ---
  248. def run_sync(
  249. is_reflect: bool,
  250. front_ring: str = None,
  251. front_coaxial: str = None,
  252. back_ring: str = None,
  253. back_coaxial: str = None,
  254. front_gray: str = None,
  255. back_gray: str = None,
  256. strict_mode: bool = True,
  257. img_index: int = None
  258. ):
  259. if os.name == 'nt':
  260. asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
  261. return asyncio.run(process_card_images(
  262. is_reflect=is_reflect,
  263. path_front_ring=front_ring,
  264. path_front_coaxial=front_coaxial,
  265. path_back_ring=back_ring,
  266. path_back_coaxial=back_coaxial,
  267. path_front_gray=front_gray,
  268. path_back_gray=back_gray,
  269. strict_mode=strict_mode,
  270. img_index=img_index
  271. ))
  272. if __name__ == "__main__":
  273. # --- 测试配置 ---
  274. IS_REFLECT = True
  275. BASE_PATH = r"C:\Code\ML\Image\Card\img20_test"
  276. # 模拟循环处理
  277. for img_num in range(1, 11):
  278. print(f">>>>> 处理图片组: {img_num}")
  279. # 构造路径 (假设文件名格式如下,可根据实际修改)
  280. p_front_coaxial = os.path.join(BASE_PATH, f"{img_num}_front_coaxial.jpg")
  281. p_front_ring = os.path.join(BASE_PATH, f"{img_num}_front_ring.jpg")
  282. p_back_coaxial = os.path.join(BASE_PATH, f"{img_num}_back_coaxial.jpg")
  283. # p_back_ring = os.path.join(BASE_PATH, f"{img_num}_back_ring_0_1.jpg")
  284. p_back_ring = r"C:\Code\ML\Image\Card\b2.jpg"
  285. # 假设有灰度图 (如果没有文件,设置为 None)
  286. # p_front_gray = os.path.join(BASE_PATH, f"{img_num}_front_gray.jpg")
  287. # if not os.path.exists(p_front_gray): p_front_gray = None # 演示用
  288. #
  289. # p_back_gray = None # 演示不传背面灰度图
  290. p_front_gray = os.path.join(BASE_PATH, f"{img_num}_front_gray.jpg")
  291. p_back_gray = os.path.join(BASE_PATH, f"{img_num}_back_gray.jpg")
  292. # 调用 (严格模式)
  293. cid = run_sync(
  294. is_reflect=IS_REFLECT,
  295. front_coaxial=p_front_coaxial,
  296. front_ring=p_front_ring,
  297. back_coaxial=p_back_coaxial,
  298. back_ring=p_back_ring,
  299. front_gray=p_front_gray,
  300. back_gray=p_back_gray,
  301. strict_mode=True, # 严格模式:缺主图会报错
  302. img_index=img_num
  303. )
  304. # 调用 (非严格模式:可以缺图)
  305. # cid = run_sync(
  306. # is_reflect=IS_REFLECT,
  307. # front_coaxial=p_front_coaxial,
  308. # # front_ring=None, # 故意不传
  309. # back_coaxial=p_back_coaxial,
  310. # # back_ring=None, # 故意不传
  311. # front_gray=p_front_gray, # 传灰度图
  312. # strict_mode=False
  313. # )
  314. if cid != -1:
  315. print(f"成功生成 ID: {cid}")
  316. else:
  317. print("生成失败")