Merge pull request #10 from KB-iGOT/main

main

Author: Haritest

app/routers/course.py

@@ -1,8 +1,8 @@
-import logging
+import time
 from fastapi import APIRouter, HTTPException
 from ..logger import logger
 from ..models import ImageResponse, ImageVariationResponse
-from ..services.course import generate_course_summary, generate_image_prompt, generate_image
+from ..services.course import generate_course_summary, generate_image_prompt, generate_image, generate_public_url
 from ..services.image_variation import generate_image_variations
 
 router = APIRouter(
@@ -13,9 +13,11 @@
 @router.get("/course/{course_id}", response_model=ImageResponse,summary= "Create thumbnail from course title, description, and table of contents")
 def generate_course_image(course_id: str):
     try:
-        final_summary = generate_course_summary(course_id)
+        content, final_summary = generate_course_summary(course_id)
         image_prompt = generate_image_prompt(final_summary)
-        image_url = generate_image(image_prompt)
+        image_data = generate_image(image_prompt)
+        image_url = generate_public_url(content, image_data)
+
         return {
             "final_summary": final_summary,
             "image_prompt": image_prompt,
@@ -28,8 +30,10 @@ def generate_course_image(course_id: str):
 @router.get("/variations/course/{course_id}", response_model=ImageVariationResponse,summary= "Generate thumbnail variations from an existing course thumbnail")
 def generate_course_image_variations(course_id: str):
     try:
+        start_time = time.time()
         logger.info(f"Course ID : {course_id}")
         image_urls = generate_image_variations(course_id)
+        print("Time took to process the request and return response is {} sec".format(time.time() - start_time))
         return ImageVariationResponse(images=image_urls)
     except Exception as e:
         logger.exception("Error while generating the image variations")

app/services/course.py

@@ -1,17 +1,26 @@
+from PIL import Image
+import base64
+import io
 import os
 from dotenv import load_dotenv
 from typing import Dict
+import urllib.parse
 import requests
 from langchain_openai import ChatOpenAI
 from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
-from langchain.chains.combine_documents import create_stuff_documents_chain
-from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
 from langchain_core.output_parsers import StrOutputParser
+import openai
+from ..utils import get_extension_from_mimetype
 
 # Load environment variables from .env file
 load_dotenv()
 KB_API_HOST = os.environ["KB_API_HOST"]
 
+from ..libs.storage import GCPStorage
+storage = GCPStorage()
+STORAGE_THUMBNAIL_FOLDER=os.environ["STORAGE_THUMBNAIL_FOLDER"]
+STORAGE_PROXY_PATH=os.environ["STORAGE_PROXY_PATH"]
+
 llm = ChatOpenAI(temperature=0, model_name="gpt-4o-mini")
 # Define prompt
 prompt = ChatPromptTemplate.from_messages(
@@ -107,7 +116,7 @@ def generate_course_summary(course_id: str):
         "description": course_details["result"]["content"]["description"], 
         "toc": formatted_toc
       })
-    return result
+    return course_details["result"]["content"], result
 
 def generate_image_prompt(summary: str):
     prompt = PromptTemplate.from_template(prompt_template)
@@ -125,12 +134,49 @@ def generate_image_prompt(summary: str):
 def generate_image(image_prompt: str):
     ############################
     # Generate a thumbnail image
-    image_url = DallEAPIWrapper(model="dall-e-3", size="1792x1024").run(f"""
+    prompt = f"""
     Do not print any text on image, just use it AS-IS:
     {image_prompt}
 
     Guidelines:
     - Please ensure that the image does not include any text or human imagery.
     - Generate image without map of india.
-    """)
-    return image_url
+    """
+    response = openai.images.generate(
+        prompt=prompt,
+        model="dall-e-3",
+        size="1024x1024",
+        quality="standard",
+        n=1,
+        response_format="b64_json"
+    )
+    return response.data[0].b64_json
+
+# Compress and convert image to JPEG
+def compress_image(image_data):
+    # Decode the base64 image data
+    image_bytes = base64.b64decode(image_data)
+    
+    # Load the image into a PIL Image object
+    image = Image.open(io.BytesIO(image_bytes))
+
+    # Convert to JPEG and compress the image
+    compressed_buffer = io.BytesIO()
+    image.save(compressed_buffer, format="JPEG", quality=80)  # Adjust quality if needed
+    compressed_image_data = compressed_buffer.getvalue()
+    return compressed_image_data
+
+def generate_public_url(content, image_data, mime_type = None):
+
+    if mime_type is None:
+       mime_type = "image/jpeg"
+
+    compressed_image_data = compress_image(image_data)
+    
+    extension = get_extension_from_mimetype(mime_type)
+    filename = f"{content["name"]}.{extension}"
+    filepath = os.path.join(STORAGE_THUMBNAIL_FOLDER, content["identifier"], filename)
+    storage.write_file(filepath, compressed_image_data, mime_type)
+    # image_urls.append(storage.public_url(filepath))
+    public_url = urllib.parse.urljoin(KB_API_HOST, os.path.join(STORAGE_PROXY_PATH, content["identifier"], filename))
+    return public_url

demo_app/README.md

@@ -0,0 +1,13 @@
+# Image Variation Generator Streamlit App
+
+## Prerequisites
+
+Before you begin, ensure you have the following installed on your system:
+
+- Thumbnail generation API
+- Streamlit library installed: `pip install streamlit`
+
+## Start the Streamlit app:
+```
+streamlit run frontend.py
+```

demo_app/frontend.py

@@ -0,0 +1,56 @@
+import requests
+import streamlit as st
+
+def fetch_image_variations(course_id):
+    """Fetches image variations for a given course ID from the API.
+
+    Args:
+        course_id (str): The ID of the course.
+
+    Returns:
+        list[str]: A list of image URLs.
+
+    Raises:
+        requests.exceptions.RequestException: If there's an error making the API request.
+    """
+
+    url = f"http://localhost:8000/v1/image/variations/course/{course_id}"
+    headers = {"accept": "application/json"}
+
+    try:
+        response = requests.get(url, headers=headers)
+        response.raise_for_status()
+        return response.json()["images"]
+    except requests.exceptions.RequestException as e:
+        raise e
+
+def display_images(images):
+    """Displays a list of image URLs in a 4-column grid.
+
+    Args:
+        images (list[str]): A list of image URLs.
+    """
+
+    cols = st.columns(4)
+    for i, image_url in enumerate(images):
+        cols[i].image(image_url)
+
+def main():
+    """Main function for the Streamlit application."""
+
+    st.title("Image Variation Generator")
+
+    course_id = st.text_input("Enter Course ID:")
+
+    generate_button = st.button("Generate Images")
+    if generate_button:
+        with st.spinner("Processing..."):
+            try:
+                images = fetch_image_variations(course_id)
+                st.success("Images generated successfully!")
+                display_images(images)
+            except requests.exceptions.RequestException as e:
+                st.error(f"Error fetching images: {e}")
+
+if __name__ == "__main__":
+    main()