Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Convert images to jpeg #2737

Merged
merged 3 commits into from
Oct 9, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 26 additions & 3 deletions backend/danswer/server/query_and_chat/chat_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import uuid
from collections.abc import Callable
from collections.abc import Generator
from typing import Tuple

from fastapi import APIRouter
from fastapi import Depends
Expand All @@ -11,6 +12,7 @@
from fastapi import Response
from fastapi import UploadFile
from fastapi.responses import StreamingResponse
from PIL import Image
from pydantic import BaseModel
from sqlalchemy.orm import Session

Expand Down Expand Up @@ -508,6 +510,21 @@ def seed_chat(
"""File upload"""


def convert_to_jpeg(file: UploadFile) -> Tuple[io.BytesIO, str]:
try:
with Image.open(file.file) as img:
if img.mode != "RGB":
img = img.convert("RGB")
jpeg_io = io.BytesIO()
img.save(jpeg_io, format="JPEG", quality=85)
jpeg_io.seek(0)
return jpeg_io, "image/jpeg"
except Exception as e:
raise HTTPException(
status_code=400, detail=f"Failed to convert image: {str(e)}"
)


@router.post("/file")
def upload_files_for_chat(
files: list[UploadFile],
Expand Down Expand Up @@ -570,19 +587,25 @@ def upload_files_for_chat(
for file in files:
if file.content_type in image_content_types:
file_type = ChatFileType.IMAGE
# Convert image to JPEG
file_content, new_content_type = convert_to_jpeg(file)
elif file.content_type in document_content_types:
file_type = ChatFileType.DOC
file_content = io.BytesIO(file.file.read())
new_content_type = file.content_type or ""
else:
file_type = ChatFileType.PLAIN_TEXT
file_content = io.BytesIO(file.file.read())
new_content_type = file.content_type or ""

# store the raw file
# store the file (now JPEG for images)
file_id = str(uuid.uuid4())
file_store.save_file(
file_name=file_id,
content=file.file,
content=file_content,
display_name=file.filename,
file_origin=FileOrigin.CHAT_UPLOAD,
file_type=file.content_type or file_type.value,
file_type=new_content_type or file_type.value,
)

# if the file is a doc, extract text and store that so we don't need
Expand Down
Loading