| 23 | |
| 24 | |
| 25 | def retrieve(class_prompt, class_data_dir, num_class_images): |
| 26 | factor = 1.5 |
| 27 | num_images = int(factor * num_class_images) |
| 28 | client = ClipClient( |
| 29 | url="https://knn.laion.ai/knn-service", indice_name="laion5B-H-14", num_images=num_images, aesthetic_weight=0.1 |
| 30 | ) |
| 31 | |
| 32 | os.makedirs(f"{class_data_dir}/images", exist_ok=True) |
| 33 | if len(list(Path(f"{class_data_dir}/images").iterdir())) >= num_class_images: |
| 34 | return |
| 35 | |
| 36 | while True: |
| 37 | class_images = client.query(text=class_prompt) |
| 38 | if len(class_images) >= factor * num_class_images or num_images > 1e4: |
| 39 | break |
| 40 | else: |
| 41 | num_images = int(factor * num_images) |
| 42 | client = ClipClient( |
| 43 | url="https://knn.laion.ai/knn-service", |
| 44 | indice_name="laion5B-H-14", |
| 45 | num_images=num_images, |
| 46 | aesthetic_weight=0.1, |
| 47 | ) |
| 48 | |
| 49 | count = 0 |
| 50 | total = 0 |
| 51 | pbar = tqdm(desc="downloading real regularization images", total=num_class_images) |
| 52 | |
| 53 | with open(f"{class_data_dir}/caption.txt", "w") as f1, open(f"{class_data_dir}/urls.txt", "w") as f2, open( |
| 54 | f"{class_data_dir}/images.txt", "w" |
| 55 | ) as f3: |
| 56 | while total < num_class_images: |
| 57 | images = class_images[count] |
| 58 | count += 1 |
| 59 | try: |
| 60 | img = requests.get(images["url"]) |
| 61 | if img.status_code == 200: |
| 62 | _ = Image.open(BytesIO(img.content)) |
| 63 | with open(f"{class_data_dir}/images/{total}.jpg", "wb") as f: |
| 64 | f.write(img.content) |
| 65 | f1.write(images["caption"] + "\n") |
| 66 | f2.write(images["url"] + "\n") |
| 67 | f3.write(f"{class_data_dir}/images/{total}.jpg" + "\n") |
| 68 | total += 1 |
| 69 | pbar.update(1) |
| 70 | else: |
| 71 | continue |
| 72 | except Exception: |
| 73 | continue |
| 74 | return |
| 75 | |
| 76 | |
| 77 | def parse_args(): |