Create helper.py
Browse files
helper.py
ADDED
@@ -0,0 +1,82 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from PIL import Image, ImageDraw
|
2 |
+
import io
|
3 |
+
import base64
|
4 |
+
import requests
|
5 |
+
import json
|
6 |
+
from typing import List, Dict, Any
|
7 |
+
from openai import OpenAI
|
8 |
+
|
9 |
+
def resize_image(image: Image.Image) -> Image.Image:
|
10 |
+
new_height = int(image.height * 512 / image.width)
|
11 |
+
return image.resize((512, new_height))
|
12 |
+
|
13 |
+
def convert_image_to_base64(image: Image.Image) -> str:
|
14 |
+
buffered = io.BytesIO()
|
15 |
+
image.save(buffered, format="PNG")
|
16 |
+
return base64.b64encode(buffered.getvalue()).decode()
|
17 |
+
|
18 |
+
def post_request_and_parse_response(url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
|
19 |
+
headers = {"Content-Type": "application/json"}
|
20 |
+
response = requests.post(url, json=payload, headers=headers)
|
21 |
+
byte_data = response.content
|
22 |
+
decoded_string = byte_data.decode("utf-8")
|
23 |
+
dict_data = json.loads(decoded_string)
|
24 |
+
return dict_data
|
25 |
+
|
26 |
+
def draw_bounding_boxes_for_textract(image: Image.Image, json_data: str) -> Image.Image:
|
27 |
+
draw = ImageDraw.Draw(image)
|
28 |
+
try:
|
29 |
+
data = json_data
|
30 |
+
blocks = json.loads(data['body']) if 'body' in data else None
|
31 |
+
except json.JSONDecodeError:
|
32 |
+
return image
|
33 |
+
|
34 |
+
if blocks is None:
|
35 |
+
return image
|
36 |
+
|
37 |
+
for item in blocks:
|
38 |
+
if 'BlockType' in item and item['BlockType'] in ['LINE', 'WORD']:
|
39 |
+
bbox = item['Geometry']['BoundingBox']
|
40 |
+
left, top, width, height = bbox['Left'], bbox['Top'], bbox['Width'], bbox['Height']
|
41 |
+
left_top = (left * image.width, top * image.height)
|
42 |
+
right_bottom = ((left + width) * image.width, (top + height) * image.height)
|
43 |
+
draw.rectangle([left_top, right_bottom], outline='red', width=2)
|
44 |
+
|
45 |
+
return image
|
46 |
+
|
47 |
+
def extract_text_from_textract_blocks(blocks: List[Dict[str, Any]]) -> str:
|
48 |
+
extracted_text = []
|
49 |
+
blocks = json.loads(blocks)
|
50 |
+
for block in blocks:
|
51 |
+
if isinstance(block, dict):
|
52 |
+
if block.get('BlockType') in ['LINE', 'WORD'] and 'Text' in block:
|
53 |
+
extracted_text.append(block['Text'])
|
54 |
+
return ' '.join(extracted_text)
|
55 |
+
|
56 |
+
class ChatGPTClient:
|
57 |
+
def __init__(self, api_key: str, protocol: str = "You are a helpful assistant.", body=None):
|
58 |
+
self.api_key = api_key
|
59 |
+
self.client = OpenAI(api_key=self.api_key)
|
60 |
+
self.protocol = protocol
|
61 |
+
self.body = body
|
62 |
+
self.history: List[Dict[str, str]] = [
|
63 |
+
{"role": "system", "content": self.protocol},
|
64 |
+
{"role": "user", "content": f"The content provided: {self.body}"}
|
65 |
+
]
|
66 |
+
|
67 |
+
def append_message(self, role: str, content: str) -> None:
|
68 |
+
if role in ["system", "user", "assistant"]:
|
69 |
+
self.history.append({"role": role, "content": content})
|
70 |
+
|
71 |
+
def generate_response(self, question: str) -> str:
|
72 |
+
try:
|
73 |
+
self.append_message("user", question)
|
74 |
+
response = self.client.chat.completions.create(
|
75 |
+
model="gpt-3.5-turbo",
|
76 |
+
messages=self.history
|
77 |
+
)
|
78 |
+
output = response.choices[0].message.content
|
79 |
+
self.append_message("assistant", output)
|
80 |
+
except Exception as e:
|
81 |
+
output = "Sorry, I couldn't get an answer for that."
|
82 |
+
return output
|