retrieve.py

# Copyright 2024 Custom Diffusion authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
importargparse
importos
fromioimportBytesIO
frompathlibimportPath

importrequests
fromclip_retrieval.clip_clientimportClipClient
fromPILimportImage
fromtqdmimporttqdm


defretrieve(class_prompt, class_data_dir, num_class_images):
factor=1.5
num_images=int(factor*num_class_images)
client=ClipClient(
url="https://knn.laion.ai/knn-service", indice_name="laion_400m", num_images=num_images, aesthetic_weight=0.1
 )

os.makedirs(f"{class_data_dir}/images", exist_ok=True)
iflen(list(Path(f"{class_data_dir}/images").iterdir())) >=num_class_images:
return

whileTrue:
class_images=client.query(text=class_prompt)
iflen(class_images) >=factor*num_class_imagesornum_images>1e4:
break
else:
num_images=int(factor*num_images)
client=ClipClient(
url="https://knn.laion.ai/knn-service",
indice_name="laion_400m",
num_images=num_images,
aesthetic_weight=0.1,
 )

count=0
total=0
pbar=tqdm(desc="downloading real regularization images", total=num_class_images)

with (
open(f"{class_data_dir}/caption.txt", "w") asf1,
open(f"{class_data_dir}/urls.txt", "w") asf2,
open(f"{class_data_dir}/images.txt", "w") asf3,
 ):
whiletotal<num_class_images:
images=class_images[count]
count+=1
try:
img=requests.get(images["url"], timeout=30)
ifimg.status_code==200:
_=Image.open(BytesIO(img.content))
withopen(f"{class_data_dir}/images/{total}.jpg", "wb") asf:
f.write(img.content)
f1.write(images["caption"] +"\n")
f2.write(images["url"] +"\n")
f3.write(f"{class_data_dir}/images/{total}.jpg"+"\n")
total+=1
pbar.update(1)
else:
continue
exceptException:
continue
return


defparse_args():
parser=argparse.ArgumentParser("", add_help=False)
parser.add_argument("--class_prompt", help="text prompt to retrieve images", required=True, type=str)
parser.add_argument("--class_data_dir", help="path to save images", required=True, type=str)
parser.add_argument("--num_class_images", help="number of images to download", default=200, type=int)
returnparser.parse_args()


if__name__=="__main__":
args=parse_args()
retrieve(args.class_prompt, args.class_data_dir, args.num_class_images)