Upload 6 files

helpers.py

@@ -0,0 +1,8 @@
+import base64, os
+
+def encode_image_to_base64(image_path: str) -> str:
+    with open(image_path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+def is_valid_image(file_name: str) -> bool:
+    return file_name.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp"))

image_functions.py

@@ -0,0 +1,82 @@
+import os
+import zipfile
+import logging
+from typing import List, Optional
+import base64
+from openai import OpenAI
+from dotenv import load_dotenv
+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):
+    with open(image_path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+def generate_image(prompt, size, quality):
+    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+    result = client.images.generate(
+        model="gpt-image-1",
+        prompt=prompt,
+        size=size,
+        quality=quality,
+    )
+
+    image_base64 = result.data[0].b64_json
+    image_bytes = base64.b64decode(image_base64)
+
+    return image_bytes

image_processor.py

@@ -0,0 +1,130 @@
+import os
+import zipfile
+import tempfile
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from image_functions import generate_image
+from prompt_service import get_prompts
+from r2_uploader import upload_image_to_r2
+from mongo_logger import create_log, update_log_status
+from helpers import encode_image_to_base64, is_valid_image
+
+
+def process_zip_and_generate_images(zip_path, category, size, quality, user_prompt):
+    try:
+        temp_dir = extract_zip_file(zip_path)
+
+        image_files = get_valid_image_files(temp_dir)
+        results = process_image_files(image_files, category, size, quality, user_prompt)
+        all_urls = [url for entry in results for url in entry["urls"]]
+
+        new_base = "https://adgen.statics.lookfinity.net/"
+
+        modified_urls = [new_base + url.rsplit('/', 1)[-1] for url in all_urls]
+        return modified_urls
+
+    except Exception as e:
+        log_error(f"Global error in zip processing: {e}")
+        return [], [f"Global error: {e}"]
+
+
+def extract_zip_file(zip_path):
+
+    temp_dir = tempfile.TemporaryDirectory()
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(temp_dir.name)
+
+    return temp_dir
+
+
+def get_valid_image_files(temp_dir):
+
+    valid_files = []
+    zip_contents = os.listdir(temp_dir.name)
+
+    for file in zip_contents:
+        file_path = os.path.join(temp_dir.name, file)
+        if is_valid_image(file) and "__MACOSX" not in file:
+            valid_files.append((file, file_path))
+
+    return valid_files
+
+
+def process_image_files(image_files, category, size, quality, user_prompt):
+    final_results = []
+    logs = []
+
+    with ThreadPoolExecutor(max_workers=5) as executor:
+        futures = []
+
+        for file_name, file_path in image_files:
+            log_id = create_log(category, file_name)
+            logs.append(f"Processing {file_name}")
+
+            futures.append(
+                executor.submit(
+                    process_single_image,
+                    file_name,
+                    file_path,
+                    category,
+                    size,
+                    quality,
+                    user_prompt,
+                    log_id,
+                    logs
+                )
+            )
+
+        for future in as_completed(futures):
+            result = future.result()
+            if result:
+                final_results.append(result)
+
+    return final_results
+
+
+def process_single_image(file_name, file_path, category, size, quality, user_prompt, log_id, logs):
+    try:
+        base64_img = encode_image_to_base64(file_path)
+        logs.append(f"Prompts generated for {file_name}")
+
+        prompts = get_prompts(base64_img, category, user_prompt)
+        logs.append(f"Generating images for {file_name}")
+
+        image_urls = generate_images_from_prompts(prompts, size, quality)
+
+        status = "completed" if image_urls else "failed"
+        message = "Completed successfully" if image_urls else "No images generated"
+        update_log_status(log_id, status, urls=image_urls, message=message)
+
+        if image_urls:
+            return {
+                "file_name": file_name,
+                "urls": image_urls
+            }
+        return None
+
+    except Exception as e:
+        logs.append(f"Processing failed for {file_name}: {e}")
+        update_log_status(log_id, "failed", urls=[], message=str(e))
+        return None
+
+
+def generate_images_from_prompts(prompts, size, quality):
+    image_urls = []
+
+    for prompt in prompts:
+        try:
+            img_url = generate_image(prompt, size, quality)
+
+            s3_url = upload_image_to_r2(img_url)
+            image_urls.append(s3_url)
+
+        except Exception as e:
+            log_error(f"Image generation failed: {e}")
+
+    return image_urls
+
+
+def log_error(message):
+    print(f"ERROR: {message}")

mongo_logger.py

@@ -0,0 +1,35 @@
+import os, uuid
+from pymongo import MongoClient
+from datetime import datetime, timezone
+
+client = MongoClient(os.getenv("MONGO_URI"))
+collection = client.adgenesis.records
+
+def create_log(category, file_name):
+    log_id = str(uuid.uuid4())
+    collection.insert_one({
+        "_id": log_id,
+        "category": category,
+        "file_name": file_name,
+        "status": "in_progress",
+        "urls": [],
+        "message": "Processing started",
+        "created_at": datetime.now(timezone.utc),
+        "updated_at": datetime.now(timezone.utc),
+        "lob": "search_arb"
+    })
+    return log_id
+
+def update_log_status(log_id, status, urls=None, message=None):
+    update = {
+        "$set": {
+            "status": status,
+            "updated_at": datetime.now(timezone.utc)
+        }
+    }
+    if urls:
+        update["$set"]["urls"] = urls
+    if message:
+        update["$set"]["message"] = message
+
+    collection.update_one({"_id": log_id}, update)

r2_uploader.py

@@ -0,0 +1,13 @@
+import boto3, os
+from uuid import uuid4
+
+def upload_image_to_r2(image_bytes):
+    s3 = boto3.client(
+        "s3",
+        endpoint_url=os.getenv("R2_ENDPOINT"),
+        aws_access_key_id=os.getenv("R2_ACCESS_KEY"),
+        aws_secret_access_key=os.getenv("R2_SECRET_KEY")
+    )
+    file_key = f"{uuid4().hex}.png"
+    s3.put_object(Bucket=os.getenv("R2_BUCKET_NAME"), Key=file_key, Body=image_bytes, ContentType="image/png")
+    return f"{os.getenv('R2_ENDPOINT')}/{os.getenv('R2_BUCKET_NAME')}/{file_key}"

requirements.txt

@@ -0,0 +1,3 @@
+openai
+pymongo
+boto3