优化AI推送代码-封装AmazonRekognitionUtil

Controller/AiController.py

@@ -5,15 +5,15 @@ import os
 import threading
 import time
 
-import boto3
 from django.views.generic.base import View
 
 from AnsjerPush.config import BASE_DIR
 from Model.models import UidPushModel, AiService, Device_Info
-from Object.ETkObject import ETkObject
 from Object.AiImageObject import ImageProcessingObject
+from Object.ETkObject import ETkObject
 from Object.ResponseObject import ResponseObject
 from Object.utils import LocalDateTimeUtil
+from Object.utils.AmazonRekognitionUtil import AmazonRekognitionUtil
 from Service.CommonService import CommonService
 from Service.EquipmentInfoService import EquipmentInfoService
 from Service.PushService import PushObject
@@ -37,7 +37,8 @@ class AiView(View):
         else:
             return response.json(414)
 
-    def identification(self, request_dict, response):
+    @staticmethod
+    def identification(request_dict, response):
         """
         ai识别推送
         @param request_dict: 请求数据
@@ -108,25 +109,14 @@ class AiView(View):
             image_row = 1  # 合并成一张图后，一行有几个小图
             image_processing_obj = ImageProcessingObject(dir_path, image_size, image_row)
             image_processing_obj.merge_images()
-            photo = open(dir_path + '.jpg', 'rb')  # 打开合成图
-
-            # 识别合成图片
-            maxLabels = 50  # 最大标签
-            minConfidence = 80  # 置信度
-
-            client = boto3.client(
-                'rekognition',
-                aws_access_key_id='AKIA2E67UIMD6JD6TN3J',
-                aws_secret_access_key='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',
-                region_name='us-east-1')
-            # doc: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/rekognition.html#Rekognition.Client.detect_labels
-            rekognition_res = client.detect_labels(
-                Image={'Bytes': photo.read()},
-                MaxLabels=maxLabels,
-                MinConfidence=minConfidence)
-            photo.close()
-            if rekognition_res['ResponseMetadata']['HTTPStatusCode'] != 200:
-                return response.json(5)
+
+            # 获取识别结果
+            aws_rekognition = AmazonRekognitionUtil()
+            with open(dir_path + '.jpg', 'rb') as f:
+                rekognition_res = aws_rekognition.detect_labels(f.read())
+
+            if not rekognition_res:
+                return response.json(0)
 
             label_dict = image_processing_obj.handle_rekognition_res(detect_group, rekognition_res)
             if not label_dict['label_list']:

Object/AiImageObject.py

@@ -5,9 +5,8 @@ import PIL.Image as Image
 from AnsjerPush.Config.aiConfig import LABEL_DICT, AI_IDENTIFICATION_TAGS_DICT
 
 
-# 图片加工类
 class ImageProcessingObject:
-
+    # 图片加工类
     def __init__(self, image_dir_path, image_size, image_row):
         self.image_dir_path = image_dir_path
         self.image_size = image_size

Object/utils/AmazonRekognitionUtil.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+"""
+@Author ： Rocky
+@Time ： 2022/12/12 14:28
+@File ：AmazonRekognitionUtil.py
+"""
+import boto3
+
+
+class AmazonRekognitionUtil:
+    # doc: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/rekognition.html
+    def __init__(self):
+        # 全部使用美东服务
+        self.region_name = 'us-east-1'
+        self.aws_access_key_id = 'AKIA2E67UIMD6JD6TN3J'
+        self.aws_secret_access_key = '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'
+
+        self.client = boto3.client(
+            'rekognition',
+            region_name=self.region_name,
+            aws_access_key_id=self.aws_access_key_id,
+            aws_secret_access_key=self.aws_secret_access_key,
+        )
+
+    def detect_labels(self, image):
+        """
+        识别图片标签
+        @param image: 图片二进制文件
+        @return: rekognition_res
+        """
+        rekognition_res = self.client.detect_labels(
+            Image={'Bytes': image},
+            MaxLabels=50,
+            MinConfidence=80
+        )
+        try:
+            assert rekognition_res['ResponseMetadata']['HTTPStatusCode'] == 200
+        except AssertionError:
+            return {}
+        return rekognition_res