image_functions.py

import os
import zipfile
import logging
from typing import List, Optional
import base64
from openai import OpenAI
from dotenv import load_dotenv
import logging

logger = logging.getLogger(__name__)

load_dotenv()

def unzip_images(zip_path: str, extract_to: Optional[str] = None) -> List[str]:

    logger = logging.getLogger(__name__)
    logger.debug(f"Starting to unzip images from {zip_path!r}")

    # Determine extraction directory
    if extract_to is None:
        extract_to = os.path.splitext(os.path.basename(zip_path))[0]
    logger.debug(f"Extraction directory set to {extract_to!r}")

    # Define supported image extensions
    image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}

    extracted_files: List[str] = []

    try:
        # Ensure the zip file exists
        if not os.path.isfile(zip_path):
            logger.error(f"Zip file does not exist: {zip_path}")
            return []

        # Create extraction directory if needed
        os.makedirs(extract_to, exist_ok=True)
        logger.debug(f"Ensured extraction directory exists: {extract_to!r}")

        # Open and iterate through the zip
        with zipfile.ZipFile(zip_path, 'r') as zf:
            for member in zf.infolist():
                filename = member.filename
                _, ext = os.path.splitext(filename.lower())

                # Only extract images
                if ext in image_extensions:
                    try:
                        zf.extract(member, extract_to)
                        file_path = os.path.join(extract_to, filename)
                        extracted_files.append(_encode_image_to_base64(file_path))
                        logger.info(f"Extracted image: {file_path}")
                    except Exception as e:
                        logger.error(f"Failed to extract {filename}: {e}")

            if not extracted_files:
                logger.warning(f"No images found in archive: {zip_path}")

        logger.debug(f"Finished unzipping. Total images extracted: {len(extracted_files)}")
        return extracted_files

    except zipfile.BadZipFile:
        logger.exception(f"Bad zip file or corrupt archive: {zip_path}")
        return []
    except Exception:
        logger.exception(f"Unexpected error while unzipping {zip_path}")
        return []

def _encode_image_to_base64(image_path):
    try:
        with open(image_path, "rb") as f:
            return base64.b64encode(f.read()).decode("utf-8")
    except Exception as e:
        logger.exception(f"Failed to base64 encode image: {image_path}")
        return ""


def generate_image(file_path, size, quality, category, sentiment, user_prompt, platform, blur, i=None):
    try:
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            logger.critical("OPENAI_API_KEY is not set.")
            raise RuntimeError("OPENAI_API_KEY is missing")
        client = OpenAI(api_key=api_key)

        if platform == "Google Display Network":
            size_messages = {
                0: "The aspect ratio of the image should be 1024x1024",
                1: "The aspect ratio of the image should be 1536x1024.",
                2: "The aspect ratio of the image should be 1024x1536.",
            }
            
            img_size = size_messages.get(i, "")
        else:
            img_size = ""
        
        with open(file_path, "rb") as img_file:
            if blur:
                background = "blurred background."
            else:
                background = " not blurred background."
            result = client.images.edit(
                model="gpt-image-1",
                prompt=(
                    f"You are a top-tier performance digital marketer and creative strategist with 15+ years of expertise in affiliate marketing.\n"
                    f"Your objective is to analyze the provided winning ad image, deconstruct its concept, visual composition, and color scheme, and generate a fresh, conversion-focused ad visual tailored for the {category} niche.\n"
                    f"The new design should convey a {sentiment} sentiment and incorporate the user instruction: \n {user_prompt}.\n If user has given multple choices or options to be include in the image so choose randomly relevant to the reference image."
                    f"Create a visually compelling ad optimized for {platform} Ads that is scroll-stopping, pattern-interrupting, and designed to drive high CTR and Conversion Rate. Utilize striking color combinations, dynamic contrast levels, and strategic layout compositions to command attention while aligning with the target audience avatar.\n"
                    f"Make sure the images should be realistic, not be stocky at all and raw which should look like they are shot from an iPhone with {background}.{img_size}"
                ),
                image=img_file,
                size=size,
                quality=quality,
            )

        image_base64 = result.data[0].b64_json
        image_bytes = base64.b64decode(image_base64)
        logger.info(f"Successfully generated image for {file_path}")
        return image_bytes

    except Exception as e:
        logger.exception(f"Failed to generate image for {file_path}: {e}")
        raise