lumaforge/dataset_curator.py

import os
import json
import urllib.request
import time
import random
from PIL import Image
from lumaforge.ollama_client import OllamaClient

class DatasetCurator:
    def __init__(self, data_dir="data/creative_dataset", ollama_client=None):
        self.data_dir = data_dir
        self.images_dir = os.path.join(data_dir, "images")
        self.metadata_path = os.path.join(data_dir, "metadata.jsonl")
        self.ollama = ollama_client or OllamaClient()
        
        # A much larger, diverse set of high-quality creative images (90 items)
        # representing different creative poster styles, characters, lighting, and environments.
        # A much larger, diverse set of high-quality creative images (90 items)
        # representing different creative poster styles, characters, lighting, and environments.
        self.source_images = [
            # Cyberpunk
            {"url": "https://images.unsplash.com/photo-1618005182384-a83a8bd57fbe?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Abstract digital painting with flowing neon pink and cyan ribbons on dark background."},
            {"url": "https://images.unsplash.com/photo-1607604276583-eef5d076aa5f?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Digital anime art illustration of futuristic city skyline with glowing neon signs at night."},
            {"url": "https://images.unsplash.com/photo-1542751371-adc38448a05e?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Gaming setup with glowing RGB keyboard, mouse, and neon futuristic background."},
            {"url": "https://images.unsplash.com/photo-1578894381163-e72c17f2d45f?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Dystopian cyberpunk city street at night with glowing purple and cyan neon signs, rain reflections."},
            {"url": "https://images.unsplash.com/photo-1548345680-f5475ea5df84?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Futuristic alleyway with high-tech wires, industrial pipes, and glowing green holographic interfaces."},
            {"url": "https://images.unsplash.com/photo-1522071820081-009f0129c71c?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Cyberpunk character with glowing augmented eyes looking at a virtual interface screen."},
            {"url": "https://images.unsplash.com/photo-1562157873-818bc0726f68?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Neon lit garage with a futuristic cyberpunk sports car parked under glowing purple lights."},
            {"url": "https://images.unsplash.com/photo-1504051771394-dd2e66b2e08f?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Crowded night market in a futuristic metropolis, glowing paper lanterns mixed with neon signs."},
            {"url": "https://images.unsplash.com/photo-1549490349-8643362247b5?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Detail shot of cybernetic hand with glowing yellow circuit lines resting on a metallic table."},
            {"url": "https://images.unsplash.com/photo-1535223289827-42f1e9919769?w=512&h=512&fit=crop", "category": "Cyberpunk", "default_caption": "Intricate server room with racks of computers emitting blinking blue and red indicator lights."},

            # Sci-Fi
            {"url": "https://images.unsplash.com/photo-1508739773434-c26b3d09e071?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "An astronaut standing on a remote desert planet looking at a massive red moon rise."},
            {"url": "https://images.unsplash.com/photo-1541185933-ef5d8ed016c2?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "A retro-futuristic rocket ship landing on Mars with volcanic dust and hazy red skies."},
            {"url": "https://images.unsplash.com/photo-1451187580459-43490279c0fa?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "Abstract view of planet Earth from deep space with glowing blue atmosphere and visible continents."},
            {"url": "https://images.unsplash.com/photo-1446776811953-b23d57bd21aa?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "A satellite solar array catching bright sunlight in outer space with stars in the background."},
            {"url": "https://images.unsplash.com/photo-1516339901601-2e1d62dc0c45?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "Mysterious deep blue space nebula cloud with shining stars and cosmic dust."},
            {"url": "https://images.unsplash.com/photo-1506703719100-a0f3a48c0f86?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "A majestic spiral galaxy swirling in empty space with clusters of bright stars and dark dust lanes."},
            {"url": "https://images.unsplash.com/photo-1614728894747-a83421e2b9c9?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "A space shuttle preparing to launch on a pad under a dramatic twilight sky."},
            {"url": "https://images.unsplash.com/photo-1581822261290-991b38693d1b?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "Retro rocket illustration on landing pad under deep starry skies."},
            {"url": "https://images.unsplash.com/photo-1502134249126-9f3755a50d78?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "A colossal space station floating in orbit above a rings-encircled gas giant planet."},
            {"url": "https://images.unsplash.com/photo-1518364538800-6bcb3f25da49?w=512&h=512&fit=crop", "category": "Sci-Fi", "default_caption": "Vintage rocket capsule floating in a swirling asteroid belt."},

            # Movie Poster
            {"url": "https://images.unsplash.com/photo-1518709268805-4e9042af9f23?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Dramatic fantasy movie poster concept of a dark medieval castle standing on a high cliff."},
            {"url": "https://images.unsplash.com/photo-1536440136628-849c177e76a1?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Cinematic film poster style scene of movie theater seats and projector beam cutting through haze."},
            {"url": "https://images.unsplash.com/photo-1489599849927-2ee91cede3ba?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Cinematic silhouette of a person standing at the end of a dark theater aisle with screen glow."},
            {"url": "https://images.unsplash.com/photo-1517604931442-7e0c8ed2963c?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Vintage cinema building exterior at night with a bright glowing sign board and retro lamps."},
            {"url": "https://images.unsplash.com/photo-1440404653325-ab127d49abc1?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Cinematic film strip frames in warm studio glowing bokeh background."},
            {"url": "https://images.unsplash.com/photo-1505686994434-e3cc5abf1330?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Cinematic film strip frames glowing with soft golden warm lights on a black background."},
            {"url": "https://images.unsplash.com/photo-1585647347483-22b66260dfff?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Suspense movie poster style scene of a foggy wooden pier leading into a dark calm lake."},
            {"url": "https://images.unsplash.com/photo-1594909122845-11baa439b7bf?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Abstract conceptual film poster of red and blue light beams intersecting on dark textured background."},
            {"url": "https://images.unsplash.com/photo-1485846234645-a62644f84728?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Cinematic shot of film director chair on set with studio lights blurred in the background."},
            {"url": "https://images.unsplash.com/photo-1478720568477-152d9b164e26?w=512&h=512&fit=crop", "category": "Movie Poster", "default_caption": "Dramatic movie poster layout of a winding mountain road at dusk under deep indigo skies."},

            # Typography
            {"url": "https://images.unsplash.com/photo-1561070791-2526d30994b5?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Bold neon pink sign showing the word 'HELLO' on brick wall."},
            {"url": "https://images.unsplash.com/photo-1509198397868-475647b2a1e5?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Vintage theater marquee letters spelling out 'CINEMA' in glowing incandescent lightbulbs."},
            {"url": "https://images.unsplash.com/photo-1563245372-f21724e3856d?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Minimalist graphic poster with big clean sans-serif typography text overlay on abstract shapes."},
            {"url": "https://images.unsplash.com/photo-1580136579312-94651dfd596d?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Abstract typographic poster with overlapping letters in black and white geometric layout."},
            {"url": "https://images.unsplash.com/photo-1525253086316-d0c936c814f8?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Clean minimalist branding poster with elegant serif text on a warm beige background."},
            {"url": "https://images.unsplash.com/photo-1579783900882-c0d3dad7b119?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Close up of old retro wooden printing press letters spelling out 'ART' in ink on paper."},
            {"url": "https://images.unsplash.com/photo-1513542789411-b6a5d4f31634?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Cursive neon light letters showing a creative phrase in soft cyan glowing on dark wall."},
            {"url": "https://images.unsplash.com/photo-1516962215378-7fa2e137ae93?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Retro style typography poster with bold 3D letters spelling 'CREATIVE' in warm gradient colors."},
            {"url": "https://images.unsplash.com/photo-1582213782179-e0d53f98f2ca?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Symmetric abstract vector graphic poster with bold sans-serif text aligned in columns."},
            {"url": "https://images.unsplash.com/photo-1504198453319-5ce911bafcde?w=512&h=512&fit=crop", "category": "Typography", "default_caption": "Chalkboard vintage lettering on a black wall, rustic sign design."},

            # Stylized Portrait
            {"url": "https://images.unsplash.com/photo-1579783902614-a3fb3927b6a5?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Classic oil painting portrait of a young woman with warm dramatic light and floral details."},
            {"url": "https://images.unsplash.com/photo-1544005313-94ddf0286df2?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Cinematic close-up portrait of a woman in low studio lighting, warm volumetric gradients."},
            {"url": "https://images.unsplash.com/photo-1500648767791-00dcc994a43e?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Black and white studio portrait of a man with high contrast key lighting, shadow depth."},
            {"url": "https://images.unsplash.com/photo-1438761681033-6461ffad8d80?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Polished studio headshot portrait of a woman looking forward, soft natural light."},
            {"url": "https://images.unsplash.com/photo-1494790108377-be9c29b29330?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Creative portrait of a smiling woman with colorful pink and yellow rim lighting in studio."},
            {"url": "https://images.unsplash.com/photo-1506794778202-cad84cf45f1d?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Cinematic portrait of a man in low light with dramatic side lighting and dark background."},
            {"url": "https://images.unsplash.com/photo-1531746020798-e6953c6e8e04?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Artistic portrait of a model with dramatic lighting and warm shadows, soft focus look."},
            {"url": "https://images.unsplash.com/photo-1554151228-14d9def656e4?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Close-up artistic portrait of a woman with vibrant makeup and soft lighting on grey backing."},
            {"url": "https://images.unsplash.com/photo-1517841905240-472988babdf9?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Studio portrait headshot of a smiling model in soft lighting."},
            {"url": "https://images.unsplash.com/photo-1508214751196-bcfd4ca60f91?w=512&h=512&fit=crop", "category": "Stylized Portrait", "default_caption": "Studio portrait of a model with colorful neon blue and magenta key lights, high contrast."},

            # Character Art
            {"url": "https://images.unsplash.com/photo-1534528741775-53994a69daeb?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "High detail studio photography portrait of a model with colored neon rim lighting."},
            {"url": "https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Close up headshot of a smiling man with textured grey background, dramatic key lighting."},
            {"url": "https://images.unsplash.com/photo-1522075469751-3a6694fb2f61?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Isolated studio character portrait of a creator with warm rim lights on grey backing."},
            {"url": "https://images.unsplash.com/photo-1566753323558-f4e0952af115?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Studio character sheet pose of a young model wearing modern athletic apparel, white background."},
            {"url": "https://images.unsplash.com/photo-1501196354995-cbb51c65aaea?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Polished portrait of a character with dramatic side lighting showing fine facial details."},
            {"url": "https://images.unsplash.com/photo-1580489944761-15a19d654956?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Studio portrait of a model wearing high-tech glasses with colorful light reflections."},
            {"url": "https://images.unsplash.com/photo-1618336753974-aae8e04506aa?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Digital painting concept art of a fantasy character with silver armor and glowing elements."},
            {"url": "https://images.unsplash.com/photo-1560250097-0b93528c311a?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Polished studio headshot of a confident executive character."},
            {"url": "https://images.unsplash.com/photo-1519085360753-af0119f7cbe7?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Business character posing confidently in studio, corporate headshot layout."},
            {"url": "https://images.unsplash.com/photo-1500048993953-d23a436266cf?w=512&h=512&fit=crop", "category": "Character Art", "default_caption": "Candid portrait headshot of a smiling programmer character."},

            # Fantasy Landscape
            {"url": "https://images.unsplash.com/photo-1506744038136-46273834b3fb?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Cinematic wide oil painting of a deep valley with rivers running through green mountains."},
            {"url": "https://images.unsplash.com/photo-1470071459604-3b5ec3a7fe05?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Mysterious fantasy forest environment with glowing mist, tall ancient mossy trees."},
            {"url": "https://images.unsplash.com/photo-1447752875215-b2761acb3c5d?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Sun rays breaking through a dense forest canopy, glowing light rays on forest floor."},
            {"url": "https://images.unsplash.com/photo-1469474968028-56623f02e42e?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Epic mountain landscape at sunset with warm orange glowing clouds and deep purple shadow peaks."},
            {"url": "https://images.unsplash.com/photo-1501785888041-af3ef285b470?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Peaceful lake scene surrounded by high alpine peaks, reflection of mountains on water surface."},
            {"url": "https://images.unsplash.com/photo-1507525428034-b723cf961d3e?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Beautiful tropical beach with golden sand, palm tree silhouettes, and glowing pink sunset."},
            {"url": "https://images.unsplash.com/photo-1501854140801-50d01698950b?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Beautiful lush green valley under massive fluffy clouds."},
            {"url": "https://images.unsplash.com/photo-1518495973542-4542c06a5843?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Magical sunlight filtering through ancient oak trees, soft mist rising from forest floor."},
            {"url": "https://images.unsplash.com/photo-1500530855697-b586d89ba3ee?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Epic deep canyon landscape with a winding blue river under a bright golden sky."},
            {"url": "https://images.unsplash.com/photo-1513836279014-a89f7a76ae86?w=512&h=512&fit=crop", "category": "Fantasy Landscape", "default_caption": "Mystic waterfalls cascading down a vertical mossy cliff inside a hidden jungle cavern."},

            # Vector Art
            {"url": "https://images.unsplash.com/photo-1531403009284-440f080d1e12?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Modern vector illustration of graphic design workspace with monitors, colorful flat icons."},
            {"url": "https://images.unsplash.com/photo-1620641788421-7a1c342ea42e?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Abstract glassmorphism vector background with frosted glass cards and vibrant mesh gradients."},
            {"url": "https://images.unsplash.com/photo-1618005198143-d5660460c55f?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Minimalist vector landscape illustration with flat geometric mountains and warm terracotta sun."},
            {"url": "https://images.unsplash.com/photo-1558655146-d09347e92766?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Cute colorful flat vector icon sheet of office elements, rounded corners, clean lines."},
            {"url": "https://images.unsplash.com/photo-1626785774573-4b799315345d?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Elegant abstract vector pattern background of thin gold geometric lines on dark navy blue background."},
            {"url": "https://images.unsplash.com/photo-1634017839464-5c339ebe3cb4?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Colorful 3D abstract vector rendering of geometric spheres and rings, neon mesh shading."},
            {"url": "https://images.unsplash.com/photo-1550684848-fac1c5b4e853?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Modern flat vector banner design with colorful abstract fluid shapes overlapping."},
            {"url": "https://images.unsplash.com/photo-1558591710-4b4a1ae0f04d?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Abstract vector pattern of clean black geometric lines on a bright white background."},
            {"url": "https://images.unsplash.com/photo-1635070041078-e363dbe005cb?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Vibrant abstract vector mesh gradient background with organic flowing shapes, blue and purple."},
            {"url": "https://images.unsplash.com/photo-1557672172-298e090bd0f1?w=512&h=512&fit=crop", "category": "Vector Art", "default_caption": "Colorful abstract fluid painting mesh background."},

            # Product Mockup
            {"url": "https://images.unsplash.com/photo-1523275335684-37898b6baf30?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Clean product mockup photograph of a white smart watch sitting on a grey concrete pedestal."},
            {"url": "https://images.unsplash.com/photo-1505740420928-5e560c06d30e?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Studio overhead photo of black wireless headphones resting on warm wooden table surface."},
            {"url": "https://images.unsplash.com/photo-1526170375885-4d8ecf77b99f?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Retro style film camera product photo with high contrast lighting, black textured backdrop."},
            {"url": "https://images.unsplash.com/photo-1542291026-7eec264c27ff?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Vibrant product shot of red athletic running shoe sitting on glowing yellow studio background."},
            {"url": "https://images.unsplash.com/photo-1583394838336-acd977736f90?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Minimalist aesthetic product display pedestal with dry palm leaf shadows on neutral beige wall."},
            {"url": "https://images.unsplash.com/photo-1572635196237-14b3f281503f?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Clean product photo of elegant retro sunglasses sitting on a white stone tile, soft key light."},
            {"url": "https://images.unsplash.com/photo-1560343090-f0409e92791a?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Studio shot of a single brown leather shoe resting on a stark white background."},
            {"url": "https://images.unsplash.com/photo-1586495777744-4413f21062fa?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Elegant cosmetic product bottle mock up resting on a concrete slab, volumetric lighting."},
            {"url": "https://images.unsplash.com/photo-1608231387042-66d1773070a5?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Minimalist mock up of white ceramic coffee mug sitting on a soft grey surface, side shadows."},
            {"url": "https://images.unsplash.com/photo-1524678606370-a47ad25cb82a?w=512&h=512&fit=crop", "category": "Product Mockup", "default_caption": "Overhead studio shot of modern white wireless earbuds inside charging case, yellow background."}
        ]

    def download_and_curate(self, limit=90, use_ollama_captioning=True) -> int:
        """
        Downloads target images, filters low-quality / duplicates,
        captions them, and writes the metadata in Hugging Face format.
        """
        print(f"[DatasetCurator] Starting curation and download process for {limit} items...")
        os.makedirs(self.images_dir, exist_ok=True)
        
        curated_count = 0
        metadata_records = []
        
        # Deduplication tracker (using simple image hash similarity)
        hashes = []
        
        # Limit cannot exceed available source images
        actual_limit = min(limit, len(self.source_images))
        selected_sources = self.source_images[:actual_limit]
        
        for idx, src in enumerate(selected_sources):
            img_filename = f"image_{idx:03d}.jpg"
            img_path = os.path.join(self.images_dir, img_filename)
            
            print(f" -> Curating item {idx+1}/{len(selected_sources)} ({src['category']})...")
            
            # Download image
            try:
                headers = {'User-Agent': 'Mozilla/5.0'}
                req = urllib.request.Request(src["url"], headers=headers)
                with urllib.request.urlopen(req, timeout=10) as response:
                    with open(img_path, 'wb') as f:
                        f.write(response.read())
            except Exception as e:
                print(f"    [Warning] Failed to download {src['url']}: {e}")
                # Fallback: create a solid color image if offline
                self._create_fallback_image(img_path, idx)
                
            # Perform Curation Step 1: Quality Checks
            if not self._check_image_validity(img_path):
                print(f"    [Curation Filter] Image {img_filename} failed quality checks. Skipping.")
                if os.path.exists(img_path):
                    os.remove(img_path)
                continue
                
            # Perform Curation Step 2: Deduplication Checks (Overfitting Control)
            img = Image.open(img_path)
            img_hash = self._calculate_average_hash(img)
            
            is_duplicate = False
            for existing_hash in hashes:
                similarity = self._hash_similarity(img_hash, existing_hash)
                if similarity > 0.82: # More strict threshold for near-duplicates
                    is_duplicate = True
                    break
                    
            if is_duplicate:
                print(f"    [Curation Filter] Image {img_filename} identified as a near-duplicate. Skipping.")
                img.close()
                os.remove(img_path)
                continue
                
            hashes.append(img_hash)
            img.close()
            
            # Perform Curation Step 3: Caption Curation
            caption = src["default_caption"]
            if use_ollama_captioning and self.ollama.check_connection():
                # Yield to let local Ollama run (simulate a slight gap to prevent hammering the local server)
                refined_caption = self._generate_ollama_caption(src["category"], caption)
                if refined_caption:
                    caption = refined_caption
                    
            relative_img_path = os.path.join("images", img_filename)
            metadata_records.append({
                "file_name": relative_img_path,
                "text": caption,
                "category": src["category"]
            })
            
            curated_count += 1
            
        # Write metadata.jsonl file
        with open(self.metadata_path, "w") as f:
            for record in metadata_records:
                f.write(json.dumps(record) + "\n")
                
        print(f"[DatasetCurator] Successfully curated {curated_count} images in '{self.data_dir}'")
        return curated_count

    def _create_fallback_image(self, path: str, idx: int):
        """Creates a solid color image as an offline fallback."""
        colors = [(200, 50, 50), (50, 200, 50), (50, 50, 200), (200, 200, 50), (200, 50, 200), (50, 200, 200)]
        color = colors[idx % len(colors)]
        img = Image.new("RGB", (512, 512), color)
        img.save(path)
        img.close()

    def _check_image_validity(self, path: str) -> bool:
        """Verifies if image is valid, uncorrupted, and meets size requirements."""
        try:
            with Image.open(path) as img:
                img.verify()
            with Image.open(path) as img:
                w, h = img.size
                if w < 256 or h < 256: # Filter low-res images
                    return False
            return True
        except Exception:
            return False

    def _calculate_average_hash(self, image: Image) -> str:
        """Calculates a simple 64-bit average hash for deduplication."""
        img_gray = image.convert("L").resize((8, 8), Image.Resampling.LANCZOS)
        pixels = list(img_gray.getdata())
        avg = sum(pixels) / 64
        bits = "".join(["1" if pixel > avg else "0" for pixel in pixels])
        return bits

    def _hash_similarity(self, hash1: str, hash2: str) -> float:
        """Calculates Hamming similarity between two binary hashes."""
        matching_bits = sum(c1 == c2 for c1, c2 in zip(hash1, hash2))
        return matching_bits / 64.0

    def _generate_ollama_caption(self, category: str, default_caption: str) -> str:
        """Uses Ollama to refine or expand image captions into detailed training prompts."""
        prompt = (
            "You are a professional image caption curator. Your job is to refine an image description "
            "to make it extremely detailed, clear, and rich for text-to-image model training. "
            "Describe subjects, colors, layout, style, and lighting. Do not use conversational text or preambles. "
            f"Image Category: {category}\n"
            f"Base description: {default_caption}\n\n"
            "Refined Caption:"
        )
        
        data = {
            "model": self.ollama.model,
            "prompt": prompt,
            "stream": False
        }
        res = self.ollama._call_api("/api/generate", data)
        if res:
            refined = res.get("response", "").strip().strip('"')
            if len(refined) > 10:
                return refined
        return None