mirrors
/
dinov2
mirror of https://github.com/facebookresearch/dinov2.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix stuff

pull/422/head
Etienne Guevel 2024-05-26 01:00:55 +02:00
parent 9d6d6d42b5
commit 279ec583f9
2 changed files with 12 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
dinov2/data/datasets/custom_image_dataset.py
View File

@ -38,22 +38,25 @@ class ImageDataset(Dataset):
return images
def _retrieve_images(self, path, is_valid=False, preserve=False, frac=1):
def _retrieve_images(self, path, is_valid=True, preserve=False, frac=1):
images_ini = len(self.preserved_images)
images = []
for root, _, files in os.walk(path):
images_dir = []
for file in files:
if file.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff')):
im = os.path.join(root, file)
if is_valid:
try:
Image.open(os.path.join(root, file))
images_dir.append(os.path.join(root, file))
with open(im, 'rb') as f:
image_data = f.read()
ImageDataDecoder(image_data).decode()
images_dir.append(im)
except OSError:
print(f"Image at path {os.path.join(root, file)} could not be opened.")
print(f"Image at path {im} could not be opened.")
else:
images_dir.append(os.path.join(root, file))
images_dir.append(im)
if preserve:
random.seed(24)
@ -61,6 +64,9 @@ class ImageDataset(Dataset):
split_index = int(len(images_dir) * frac)
self.preserved_images.extend(images_dir[:split_index])
images.extend(images_dir[split_index:])
else:
images.extend(images_dir)
images_end = len(self.preserved_images)
if preserve:

2
dinov2/train/train.py
View File

@ -200,7 +200,7 @@ def do_train(cfg, model, resume=False):
# save the preserved images
if dataset.preserved_images:
write_list(os.path.join(cfg.train.output_dir, 'preserved_images.pkl'))
write_list(os.path.join(cfg.train.output_dir, 'preserved_images.pkl'), dataset.preserved_images)
# sampler_type = SamplerType.INFINITE
sampler_type = SamplerType.SHARDED_INFINITE # define the sampler to use for fsdp