creek

  1"""
  2Testing Data Generation Template
  3Original file: https://colab.research.google.com/drive/1iDu3v3m3iOKyRs6IggnOPbN6tKKa9SRV
  4"""
  5
  6## Setup Environment
  7
  8!pip install requests tqdm icrawler
  9
 10import numpy as np
 11import pandas as pd
 12import os
 13import shutil
 14import urllib3
 15import logging
 16from icrawler import ImageDownloader
 17from icrawler.builtin import BingImageCrawler
 18
 19OUTPUT_ROOT = "Feature_Dataset_HQ"
 20TARGET_COUNT = 100
 21
 22urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 23
 24
 25class No403Filter(logging.Filter):
 26    def filter(self, record):
 27        return "403" not in record.getMessage()
 28
 29
 30logging.getLogger("downloader").addFilter(No403Filter())
 31
 32
 33class RobustDownloader(ImageDownloader):
 34    def __init__(self, *args, **kwargs):
 35        super().__init__(*args, **kwargs)
 36        self.session.verify = False
 37        self.session.headers.update(
 38            {
 39                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
 40                "Referer": "https://www.bing.com/",
 41            }
 42        )
 43
 44
 45SEARCH_MAPPING = {
 46    "soft-light": "soft light, gentle illumination, low contrast, wrap-around light",
 47    "specular-highlights": "strong specular highlights, shiny skin, glossy reflections, bright hotspots",
 48    "backlit": "backlit lighting from behind the subject, rim light, silhouette effect, sun flare",
 49    "studio-lighting": "professional studio lighting, controlled illumination, three-point lighting setup",
 50    "flat-lighting": "flat lighting, low contrast, even illumination, no shadows, cloudy day look",
 51    "dramatic-contrast": "dramatic high-contrast lighting, chiaroscuro, hard shadows, noir style",
 52    "diffused": "soft diffused lighting, cloudy sky light, light box look, scattered light, no hard edges",
 53}
 54
 55
 56def count_files(folder):
 57    if not os.path.exists(folder):
 58        return 0
 59    return len([f for f in os.listdir(folder) if f.endswith((".jpg", ".png", ".jpeg"))])
 60
 61
 62def trim_excess(folder, limit):
 63    files = sorted(
 64        [f for f in os.listdir(folder) if f.endswith((".jpg", ".png", ".jpeg"))]
 65    )
 66    if len(files) > limit:
 67        print(f"   -> Overshot! Trimming {len(files) - limit} extra images...")
 68        for f in files[limit:]:
 69            try:
 70                os.remove(os.path.join(folder, f))
 71            except:
 72                pass
 73
 74
 75if __name__ == "__main__":
 76    if not os.path.exists(OUTPUT_ROOT):
 77        os.makedirs(OUTPUT_ROOT)
 78
 79    print(f"Starting Recursive Download (Target: {TARGET_COUNT})...")
 80
 81    for category, query in SEARCH_MAPPING.items():
 82        print(f"\n--- Processing: {category} ---")
 83        save_dir = os.path.join(OUTPUT_ROOT, category)
 84
 85        if not os.path.exists(save_dir):
 86            os.makedirs(save_dir)
 87
 88        # Initial Search Depth
 89        search_depth = 200
 90
 91        while True:
 92            current_count = count_files(save_dir)
 93
 94            # 1. CHECK CONDITION
 95            if current_count >= TARGET_COUNT:
 96                print(f"   -> Success! Found {current_count} images.")
 97                break
 98
 99            # 2. INCREASE DEPTH
100            print(
101                f"   -> Have {current_count}/{TARGET_COUNT}. Increasing search depth to {search_depth}..."
102            )
103
104            crawler = BingImageCrawler(
105                downloader_cls=RobustDownloader,
106                feeder_threads=1,
107                parser_threads=2,
108                downloader_threads=4,
109                storage={"root_dir": save_dir},
110                log_level="ERROR",
111            )
112
113            # 3. SAFETY OFFSET
114            if current_count == 0:
115                offset_param = 0
116            else:
117                offset_param = "auto"
118
119            crawler.crawl(
120                keyword=query,
121                max_num=search_depth,
122                filters={"type": "photo"},
123                file_idx_offset=offset_param,
124            )
125
126            # Increase the search buffer for the next loop
127            search_depth += 500
128
129            # Safety break
130            if search_depth > 5000:
131                print("   -> Hit safety limit (5000). Stopping.")
132                break
133
134        # Final Trim
135        trim_excess(save_dir, TARGET_COUNT)
136        print(f"   Final Count: {count_files(save_dir)}")
137
138    print(f"\n Material Look Dataset (HQ) Download Complete!")
139
140
141## Compress dataset folder into .zip
142
143import shutil
144import os
145
146FOLDER_TO_ZIP = "MaterialLook_Dataset_HQ"
147OUTPUT_FILENAME = "MaterialLook_Dataset"  # .zip will be added automatically
148
149if __name__ == "__main__":
150    if os.path.exists(FOLDER_TO_ZIP):
151        print(f"Compressing '{FOLDER_TO_ZIP}'...")
152        shutil.make_archive(OUTPUT_FILENAME, "zip", FOLDER_TO_ZIP)
153        zip_size = os.path.getsize(f"{OUTPUT_FILENAME}.zip") / (1024 * 1024)
154        print(f" Success! Created {OUTPUT_FILENAME}.zip ({zip_size:.2f} MB)")
155    else:
156        print(f" Error: Folder '{FOLDER_TO_ZIP}' not found.")