Skip to content

Commit 488c43e

Browse files
committed
improve: enhance image format detection in process_image method
- Add proper MIME type detection for both data URLs and HTTP requests - Extract media type from Content-Type header or MIME type - Make format detection more robust and generic - Remove hardcoded PNG/JPEG format assumptions
1 parent 3270627 commit 488c43e

File tree

1 file changed

+13
-6
lines changed

1 file changed

+13
-6
lines changed

examples/pipelines/providers/aws_bedrock_claude_pipeline.py

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -227,17 +227,24 @@ def pipe(
227227

228228
def process_image(self, image: str):
229229
img_stream = None
230+
content_type = None
231+
230232
if image["url"].startswith("data:image"):
231-
if ',' in image["url"]:
232-
base64_string = image["url"].split(',')[1]
233+
mime_type, base64_string = image["url"].split(",", 1)
234+
content_type = mime_type.split(":")[1].split(";")[0]
233235
image_data = base64.b64decode(base64_string)
234-
235236
img_stream = BytesIO(image_data)
236237
else:
237-
img_stream = requests.get(image["url"]).content
238+
response = requests.get(image["url"])
239+
img_stream = BytesIO(response.content)
240+
content_type = response.headers.get('Content-Type', 'image/jpeg')
241+
242+
media_type = content_type.split('/')[-1] if '/' in content_type else content_type
238243
return {
239-
"image": {"format": "png" if image["url"].endswith(".png") else "jpeg",
240-
"source": {"bytes": img_stream.read()}}
244+
"image": {
245+
"format": media_type,
246+
"source": {"bytes": img_stream.read()}
247+
}
241248
}
242249

243250
def stream_response(self, model_id: str, payload: dict) -> Generator:

0 commit comments

Comments
 (0)