gamingflexer commited on
Commit
620e74a
·
1 Parent(s): a59c27e

Add OpenAIVision class for image description generation

Browse files
Files changed (1) hide show
  1. src/module/llm_vision.py +59 -0
src/module/llm_vision.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import base64
2
+ import requests
3
+ from config import OPENAI_API
4
+ import os
5
+
6
+ OPENAI_API = os.getenv("OPENAI_API")
7
+
8
+ """
9
+ openai_vision = OpenAIVision(api_key)
10
+ image_path = "path_to_your_image.jpg"
11
+ prompt = ""
12
+ response = openai_vision.get_image_description(prompt,image_path)
13
+ """
14
+
15
+ class OpenAIVision:
16
+ def __init__(self):
17
+ self.api_key = OPENAI_API
18
+ self.base_url = "https://api.openai.com/v1/chat/completions"
19
+
20
+ def __encode_image(self, image_path):
21
+ with open(image_path, "rb") as image_file:
22
+ return base64.b64encode(image_file.read()).decode('utf-8')
23
+
24
+ def get_image_description(self, image_path, prompt):
25
+ base64_image = self.__encode_image(image_path)
26
+
27
+ headers = {
28
+ "Content-Type": "application/json",
29
+ "Authorization": f"Bearer {self.api_key}"
30
+ }
31
+
32
+ payload = {
33
+ "model": "gpt-4-vision-preview",
34
+ "temperature": 0.0,
35
+ "messages": [
36
+ {
37
+ "role": "user",
38
+ "content": [
39
+ {
40
+ "type": "text",
41
+ "text": prompt,
42
+ },
43
+ {
44
+ "type": "image_url",
45
+ "image_url": {
46
+ "url": f"data:image/jpeg;base64,{base64_image}"
47
+ }
48
+ }
49
+
50
+ ]
51
+
52
+ }
53
+ ],
54
+ "max_tokens": 1000,
55
+
56
+ }
57
+
58
+ response = requests.post(self.base_url, headers=headers, json=payload)
59
+ return response.json()