Create image_fetcher.py

image_fetcher.py

@@ -0,0 +1,71 @@
+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)
+
+# 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:
+                    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)
+