File size: 2,459 Bytes
8b0cadb 31d1e0a 8b0cadb 31d1e0a 8b0cadb 31d1e0a 8b0cadb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
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
|