Add support for images

Author: punitarani

app.py

@@ -4,7 +4,7 @@
 from langchain.schema import Document
 
 from rabbithole import summarize_document
-from rabbithole.loader import load_file
+from rabbithole.loader import load_file, SUPPORTED_IMG_FILE_TYPES
 from rabbithole.mp3 import SUPPORTED_AV_FILE_TYPES
 
 # Global variables
@@ -36,7 +36,8 @@ def run_summarization(document: list[Document], doc_name: str):
 
 st.title("RabbitHole")
 
-uploaded_files = st.file_uploader("Upload content", type=["docx", "pdf", "txt", *SUPPORTED_AV_FILE_TYPES],
+uploaded_files = st.file_uploader("Upload content",
+                                  type=["docx", "pdf", "txt", *SUPPORTED_IMG_FILE_TYPES, *SUPPORTED_AV_FILE_TYPES],
                                   accept_multiple_files=True)
 
 if st.button("Summarize"):

rabbithole/loader.py

@@ -1,14 +1,16 @@
 """rabbithole.loader module"""
 import tempfile
 
-from langchain.document_loaders import Docx2txtLoader, PyMuPDFLoader, TextLoader
+from langchain.document_loaders import Docx2txtLoader, PyMuPDFLoader, TextLoader, UnstructuredImageLoader
 from langchain.schema import Document
 from langchain.text_splitter import CharacterTextSplitter
 from streamlit.runtime.uploaded_file_manager import UploadedFile
 
 from rabbithole.mp3 import SUPPORTED_AV_FILE_TYPES, convert_to_mp3
 from rabbithole.transcribe import transcribe
 
+SUPPORTED_IMG_FILE_TYPES = (".jpg", ".jpeg", ".png")
+
 
 def save_to_temp_file(file: UploadedFile) -> str:
     """
@@ -46,6 +48,11 @@ def load_file(file: UploadedFile) -> list[Document]:
         temp_file = save_to_temp_file(file)
         return TextLoader(file_path=temp_file).load_and_split(text_splitter=text_splitter)
 
+    # Handle image files
+    elif file.name.endswith(SUPPORTED_IMG_FILE_TYPES):
+        temp_file = save_to_temp_file(file)
+        return UnstructuredImageLoader(file_path=temp_file).load_and_split(text_splitter=text_splitter)
+
     # Handle Audio and Video files
     elif file.name.endswith(SUPPORTED_AV_FILE_TYPES):
         temp_file = save_to_temp_file(file)