akshansh36 commited on
Commit
a43b00d
·
verified ·
1 Parent(s): 1fd56ea

Update helper/process_image.py

Browse files
Files changed (1) hide show
  1. helper/process_image.py +91 -91
helper/process_image.py CHANGED
@@ -1,91 +1,91 @@
1
- import base64
2
- import httpx
3
- from pymongo import MongoClient
4
- from langchain_google_genai import ChatGoogleGenerativeAI
5
- from langchain_core.messages import HumanMessage
6
- import os
7
- import re
8
- import json
9
- from dotenv import load_dotenv
10
- load_dotenv()
11
- MONGO_URI = os.getenv("MONGO_URI")
12
- DB_NAME = os.getenv("DB_NAME")
13
- COLLECTION_NAME = os.getenv("COLLECTION_NAME")
14
- FLASH_API = os.getenv("FLASH_API")
15
- mongo_client = MongoClient(MONGO_URI)
16
- db = mongo_client[DB_NAME]
17
- collection = db[COLLECTION_NAME]
18
- collection2=db['about_company']
19
- model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0, max_tokens=None, google_api_key=FLASH_API)
20
-
21
- about_company_doc=collection2.find_one({"type":"about_company"})
22
- if about_company_doc:
23
- about_company=about_company_doc.get('company_description','')
24
-
25
- system_prompt_text = f"""Given is an image related to a company. Your task is to analyze the image, identify any text or notable visual elements, and provide a comprehensive, direct description of the image's contents, focusing on what it represents without abstract language or additional commentary. The response must be concise and focused, using only descriptive nouns and adjectives. If the image cannot be clearly described, respond with 'None.'
26
- Company information is given below to understand the context.
27
- - About Company: {about_company}
28
-
29
- Expected Output format : {{'description':'String'}}
30
- """
31
-
32
-
33
-
34
-
35
- def process_image_using_llm(image_url):
36
- try:
37
- # Download and encode the image
38
- image_data = base64.b64encode(httpx.get(image_url).content).decode("utf-8")
39
-
40
- # Create the message with a system prompt and image
41
- message = HumanMessage(
42
- content=[
43
- {"type": "text", "text": system_prompt_text},
44
- {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}},
45
- ],
46
- )
47
-
48
- # Invoke the model to extract information from the image
49
- response = model.invoke([message])
50
- print(f"llm response for {image_url} is : {response}")
51
- # Use regex to extract JSON part of the response
52
- match = re.search(r"\{.*\}", response.content.strip())
53
- if match:
54
- json_data = match.group(0) # Extract JSON-like content as a string
55
- json_data = json_data.replace("'", '"')
56
- data = json.loads(json_data) # Load as JSON
57
-
58
- # Get the description from the JSON data
59
- description = data.get("description", "None").strip()
60
-
61
- # Check if the description is "None"
62
- if description == "None":
63
- # Update MongoDB with can_find_description as False
64
- collection.update_one(
65
- {"object_url": image_url},
66
- {"$set": {"can_find_description": False}}
67
- )
68
- print(f"Marked {image_url} as can_find_description: False")
69
- return False
70
- else:
71
- # Update MongoDB with the description and set can_find_description to True
72
- collection.update_one(
73
- {"object_url": image_url},
74
- {"$set": {"description": description, "can_find_description": True}}
75
- )
76
- print("Saved description to MongoDB")
77
- return True
78
- else:
79
- print(f"No valid JSON found in the response for {image_url}. Marking as can_find_description: False")
80
- collection.update_one(
81
- {"object_url": image_url},
82
- {"$set": {"can_find_description": False}}
83
- )
84
-
85
- return False
86
- except Exception as e:
87
- print(f"Error processing {image_url}: {e}")
88
- return False
89
-
90
-
91
-
 
1
+ import base64
2
+ import httpx
3
+ from pymongo import MongoClient
4
+ from langchain_google_genai import ChatGoogleGenerativeAI
5
+ from langchain_core.messages import HumanMessage
6
+ import os
7
+ import re
8
+ import json
9
+ from dotenv import load_dotenv
10
+ load_dotenv()
11
+ MONGO_URI = os.getenv("MONGO_URI")
12
+ DB_NAME = os.getenv("DB_NAME")
13
+ COLLECTION_NAME = os.getenv("COLLECTION_NAME")
14
+ FLASH_API = os.getenv("FLASH_API")
15
+ mongo_client = MongoClient(MONGO_URI)
16
+ db = mongo_client[DB_NAME]
17
+ collection = db[COLLECTION_NAME]
18
+ collection2=db['about_company']
19
+ model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0, max_tokens=None, google_api_key=FLASH_API)
20
+
21
+ about_company_doc=collection2.find_one({"type":"about_company"})
22
+ if about_company_doc:
23
+ about_company=about_company_doc.get('company_description','')
24
+
25
+ system_prompt_text = f"""Given is an image related to a company. Your task is to analyze the image, identify any text or notable visual elements, and provide a comprehensive, direct description of the image's contents, focusing on what it represents without abstract language or additional commentary. The response must be concise and focused, using only descriptive nouns and adjectives. If the image cannot be clearly described, respond with 'None.'
26
+ Company information is given below to understand the context.
27
+ - About Company: {about_company}
28
+
29
+ Expected Output format : {{"description":"String"}}
30
+ """
31
+
32
+
33
+
34
+
35
+ def process_image_using_llm(image_url):
36
+ try:
37
+ # Download and encode the image
38
+ image_data = base64.b64encode(httpx.get(image_url).content).decode("utf-8")
39
+
40
+ # Create the message with a system prompt and image
41
+ message = HumanMessage(
42
+ content=[
43
+ {"type": "text", "text": system_prompt_text},
44
+ {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}},
45
+ ],
46
+ )
47
+
48
+ # Invoke the model to extract information from the image
49
+ response = model.invoke([message])
50
+ print(f"llm response for {image_url} is : {response}")
51
+ # Use regex to extract JSON part of the response
52
+ match = re.search(r"\{.*\}", response.content.strip())
53
+ if match:
54
+ json_data = match.group(0) # Extract JSON-like content as a string
55
+ json_data = json_data.replace("'", '"')
56
+ data = json.loads(json_data) # Load as JSON
57
+
58
+ # Get the description from the JSON data
59
+ description = data.get("description", "None").strip()
60
+
61
+ # Check if the description is "None"
62
+ if description == "None":
63
+ # Update MongoDB with can_find_description as False
64
+ collection.update_one(
65
+ {"object_url": image_url},
66
+ {"$set": {"can_find_description": False}}
67
+ )
68
+ print(f"Marked {image_url} as can_find_description: False")
69
+ return False
70
+ else:
71
+ # Update MongoDB with the description and set can_find_description to True
72
+ collection.update_one(
73
+ {"object_url": image_url},
74
+ {"$set": {"description": description, "can_find_description": True}}
75
+ )
76
+ print("Saved description to MongoDB")
77
+ return True
78
+ else:
79
+ print(f"No valid JSON found in the response for {image_url}. Marking as can_find_description: False")
80
+ collection.update_one(
81
+ {"object_url": image_url},
82
+ {"$set": {"can_find_description": False}}
83
+ )
84
+
85
+ return False
86
+ except Exception as e:
87
+ print(f"Error processing {image_url}: {e}")
88
+ return False
89
+
90
+
91
+