add: title and description generator

requirements.txt

@@ -78,4 +78,5 @@ uvloop==0.19.0
 watchfiles==0.22.0
 websockets==12.0
 replicate
+google-generativeai
 -e .

src/__init__.py

@@ -0,0 +1,5 @@
+"""
+project @ NTO-TCP-HF
+created @ 2024-10-28
+author  @ github/ishworrsubedii
+"""

src/api/image_prep_api.py

@@ -17,13 +17,17 @@ from fastapi.responses import JSONResponse
 
 from src.components.auto_crop import crop_transparent_image
 from src.components.color_extraction import ColorExtractionRMBG
+from src.components.title_des_gen import NecklaceProductListing
 
 preprocessing_router = APIRouter()
 
 rmbg: str = os.getenv("RMBG")
 
 enhancer: str = os.getenv("ENHANCER")
+prod_listing_api_key: str = os.getenv("PROD_LISTING_API_KEY")
+
 color_extraction_rmbg = ColorExtractionRMBG()
+product_listing_obj = NecklaceProductListing(prod_listing_api_key)
 
 
 def replicate_bg(input):
@@ -187,3 +191,16 @@ async def remove_background_color_extraction(image: UploadFile = File(...), hex_
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed to process image: {e}")
+
+
+@preprocessing_router.post("/title_description_generator")
+async def product_title_description_generator(image: UploadFile = File(...)):
+    image_bytes = await image.read()
+    image = Image.open(BytesIO(image_bytes)).convert("RGB")
+    try:
+        result = product_listing_obj.gen_title_desc(image=image)
+
+        return JSONResponse(content=result, status_code=200)
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to process image: {e}")

src/components/title_des_gen.py

@@ -0,0 +1,43 @@
+"""
+project @ NTO-TCP-HF
+created @ 2024-10-29
+author  @ github.com/ishworrsubedii
+"""
+import google.generativeai as genai
+from tempfile import NamedTemporaryFile
+
+
+class NecklaceProductListing:
+    def __init__(self, api_key):
+        genai.configure(api_key=api_key)
+        self.model = genai.GenerativeModel(model_name="gemini-1.5-flash")
+        self.prompt = """Analyze this necklace image and create an e-commerce product listing with:
+        1. An SEO-friendly, compelling product title (3-8 words)
+        2. A persuasive product description for online shoppers (80-120 words)
+        Format your response exactly like this:
+        Title: [Product Title]
+        Description: [Product Description]
+        For the description, include:
+        - Opening hook about the piece's beauty or uniqueness
+        - Key materials and specifications (metal type, gemstones, length, closure type)
+        - Highlight 2-3 standout design features
+        - Styling suggestions (what to wear it with, occasions)
+        - Quality/craftsmanship mentions
+        - One emotional benefit (how it makes the wearer feel)
+        Write in a professional yet warm tone that appeals to online jewelry shoppers. Focus on benefits and value proposition."""
+
+    def save_image_tempfile(self, pil_image):
+        with NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+            pil_image.save(temp_file, format="PNG")
+            return temp_file.name
+
+    def image_upload(self, image_path):
+        sample_file = genai.upload_file(path=image_path, display_name="necklace_image")
+        return sample_file
+
+    def gen_title_desc(self, image):
+        image_path = self.save_image_tempfile(image)
+        sample_file = self.image_upload(image_path)
+
+        response = self.model.generate_content([sample_file, self.prompt])
+        return response.text