增加图像分析功能

agent/agent.py

@@ -50,7 +50,15 @@ class Agent:
         """
         response = self.glm.text_response(prompt)
         response = response.content.replace(' ', '').replace('\n', '').replace('\t', '').replace(":", ": ")
-        print(response)
+        return response
+    
+    def image_logo_judgement(self, logo_path, image_url):
+        """判断图像中是否有指定品牌的logo"""
+        self.glm.set_modelname("glm-4v-plus-0111")
+        prompt = KeyWordPrompt.IMAGE_LOGO_JUDGEMENT
+        response = self.glm.multi_epoch_image_response(logo_path, image_url, prompt)
+        response = response.content
+        
         return response
 
 if __name__ == "__main__":

agent/config.py

@@ -39,6 +39,27 @@ class KeyWordPrompt:
         
     """
     
+    IMAGE_LOGO_JUDGEMENT = f"""
+        你是一个产品图像分析助手，你的任务是判断第二张图像中的产品上是否包含logo，并与第一张logo图像做对比，判断第二张图像中的logo是否与第一张的一致。
+        判断思路：
+        1. 首先判断第二张图像中是否包含logo。
+        2. 如果包含logo的话，再进行判断，是否与第一张的logo一致。
+        3. 如果与第一张的logo不一致，请根据你的经验判断图像中的logo是什么品牌，如果不知道返回'未知'
+        最终结果返回为以下给出的json格式
+        输出结果示例：
+        
+        {{
+            "is_contain_logo": false
+        }}
+        
+        
+        {{
+            "is_contain_logo": true,
+            "is_jugement_logo": true,
+            "brand_name": "李宁"
+        }}
+    """
+    
     
     
 if __name__ == "__main__":

agent/glm.py

@@ -53,11 +53,11 @@ class Glm:
         
         return response.choices[0].message
     
-    def multi_epoch_image_response(self, logo_path):
+    def multi_epoch_image_response(self, logo_path, image_url, query):
         """多轮图像问答"""
         image_base = image_to_base(logo_path)
         response = self.client.chat.completions.create(
-            model="glm-4v-plus-0111",
+            model=self.model_name,
             messages=[
                 {
                     "content": [
@@ -75,17 +75,20 @@ class Glm:
                 {
                     "content": [
                         {
-                            "image_url": {"url": "http://h2.appsimg.com/a.appsimg.com/upload/merchandise/pdcvis/2023/04/14/79/d75fa3db-3bec-45c4-b7a6-54332d42e373.jpg"},
+                            "image_url": {"url": image_url},
                             "type": "image_url"
                         },
                         {
-                            "text": "第二张图像中的产品是否包含logo，如果包含的话是否是第一张的logo？",
+                            "text": query,
                             "type": "text"
                         }
                     ],
                     "role": "user"
                 }
-            ]
+            ],
+            response_format = {
+                "type": "joson_object"
+            }
         )
         
         return response.choices[0].message
@@ -95,5 +98,26 @@ class Glm:
     
 if __name__ == '__main__':
     glm = Glm()
-    response = glm.multi_epoch_image_response("./logo/lining.jpg")
-    print(response)
+    glm.set_modelname("glm-4v-plus-0111")
+    query = f"""
+        你是一个产品图像分析助手，你的任务是判断第二张图像中的产品上是否包含logo，并与第一张logo图像做对比，判断第二张图像中的logo是否与第一张的一致。
+        判断思路：
+        1. 首先判断第二张图像中是否包含logo。
+        2. 如果包含logo的话，再进行判断，是否与第一张的logo一致。
+        3. 如果与第一张的logo不一致，请根据你的经验判断图像中的logo是什么品牌，如果不知道返回'未知'
+        最终结果返回为以下给出的json格式
+        输出结果示例：
+        
+        {{
+            "is_contain_logo": false
+        }}
+        
+        
+        {{
+            "is_contain_logo": true,
+            "is_jugement_logo": true,
+            "brand_name": "李宁"
+        }}
+    """
+    response = glm.multi_epoch_image_response("./logo/lining.jpg", "http://h2.appsimg.com/a.appsimg.com/upload/merchandise/pdcvis/613214/2024/1120/88/9b9027dd-95b7-4024-b71e-fb7cbfde16a1.jpg", query)
+    print(response.content)

webui.py

@@ -74,27 +74,36 @@ def get_license_list():
 
 def check_infringement(title, brandname):
     """ 模拟侵权检测逻辑 """
+    record = products_dao.get_one_record_by_query({"title": title})
+    image_url = record["image"][0]
+    
     if brandname not in title:
-        record = products_dao.get_one_record_by_query({"title": title})
         actual_brandname = record["brandName"]
         if actual_brandname not in brandname:
             key_word_falg = True
         else:
             key_word_falg = False
-        
-    key_word_judgement = json.loads(agent.brand_key_word_judgement(brandname, title))
     
-    res = agent.license_product_judgement(title, license_list_str)
-    print(res)
-    license_judgement = json.loads(res)
+    key_word_judgement = json.loads(agent.brand_key_word_judgement(brandname, title))
+    license_judgement = json.loads(agent.license_product_judgement(title, license_list_str))
+    logo_judgement = json.loads(agent.image_logo_judgement("./logo/lining.jpg", image_url))
+   
     key_word_falg = key_word_judgement["key_word_flag"]
     license_judgement_flag = license_judgement["in_list"]
     # license_flag = license_judgement["in_list"]
-    print(license_judgement)
     result = f"""
         关键词引流: {key_word_falg}
         是否为授权生产产品: {license_judgement_flag}
+        
+        产品LOGO图像判定:
+        图像中是否包含logo: {logo_judgement["is_contain_logo"]}
     """
+    
+    if logo_judgement["is_contain_logo"]:
+        result += 
+        f"""是否是指定品牌LOGO: {logo_judgement["is_jugement_logo"]}
+        LOGO名称: {logo_judgement["brand_name"]}
+        """
     return result
 
 def search_by_cust_id(cust_id):