image_fetcher.py

import os
import asyncio
import aiohttp
import random
from PIL import Image
from io import BytesIO
from duckduckgo_search import DDGS
import time

# Setup /tmp directory for Hugging Face
IMAGE_DIR = "/tmp/images"
os.makedirs(IMAGE_DIR, exist_ok=True)
image_nlst=[]

# Headers
def get_headers():
    user_agents = [
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
        "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:90.0)"
    ]
    return {"User-Agent": random.choice(user_agents)}

# Image validation
def is_valid_image(image):
    width, height = image.size
    aspect_ratio = round(width / height, 2)
    return width >= 854 and height >= 480 and abs(aspect_ratio - (16 / 9)) <= 0.2

# Async fetch image
async def fetch_image(session, url, name):
    try:
        async with session.get(url, timeout=10) as response:
            content = await response.read()
            image = Image.open(BytesIO(content)).convert("RGB")

            if not is_valid_image(image):
                return f"Skipped (invalid): {name}"

            unique_name = f"{name}_{int(time.time() * 1000)}.jpg"
            path = os.path.join(IMAGE_DIR, unique_name)
            image.save(path)
            return f"Saved: {unique_name}"

    except Exception as e:
        return f"Error: {name} | {e}"

# Async search + download with DDGS inside thread
async def search_and_download(session, prompt, sem):
    async with sem:
        name = prompt.replace(" ", "_").lower()
        try:
            loop = asyncio.get_event_loop()
            results = await loop.run_in_executor(None, lambda: list(DDGS().images(prompt, max_results=15)))
            for item in results:
                url = item.get("image")
                result = await fetch_image(session, url, name)
                if "Saved" in result:
                    image_nlst.append(result.lstrip("Saved:"))
                    return result
            return f"No valid image for: {prompt}"
        except Exception as e:
            return f"Search failed for {prompt}: {e}"

# Main runner
async def main(prompts):
    sem = asyncio.Semaphore(5)  # Limit concurrency
    async with aiohttp.ClientSession(headers=get_headers()) as session:
        tasks = [search_and_download(session, prompt, sem) for prompt in prompts]
        results = await asyncio.gather(*tasks)
        for res in results:
            print(res)
    return image_nlst