Spaces:
Sleeping
Sleeping
File size: 2,391 Bytes
a479aa5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import pickle
import os
import glob
dataroot = "/media/mmlab/Volume2/TrueFake"
dataset_structure = [
'Fake/FLUX.1/animals',
'Fake/FLUX.1/faces',
'Fake/FLUX.1/general',
'Fake/FLUX.1/landscapes',
'Fake/StableDiffusion1.5/animals',
'Fake/StableDiffusion1.5/faces',
'Fake/StableDiffusion1.5/general',
'Fake/StableDiffusion1.5/landscapes',
'Fake/StableDiffusion2/animals',
'Fake/StableDiffusion2/faces',
'Fake/StableDiffusion2/general',
'Fake/StableDiffusion2/landscapes',
'Fake/StableDiffusion3/animals',
'Fake/StableDiffusion3/faces',
'Fake/StableDiffusion3/general',
'Fake/StableDiffusion3/landscapes',
'Fake/StableDiffusionXL/animals',
'Fake/StableDiffusionXL/faces',
'Fake/StableDiffusionXL/general',
'Fake/StableDiffusionXL/landscapes',
'Fake/StyleGAN/images-psi-0.5',
'Fake/StyleGAN/images-psi-0.7',
'Fake/StyleGAN2/conf-f-psi-0.5',
'Fake/StyleGAN2/conf-f-psi-1',
'Fake/StyleGAN3/conf-t-psi-0.5',
'Fake/StyleGAN3/conf-t-psi-0.7',
'Real/FFHQ',
'Real/FORLAB'
]
if __name__ == "__main__":
with open("./classes_nosocial.pkl", "rb") as f:
results = pickle.load(f)
results_shared = {}
for social in ['Facebook', 'Telegram', 'Twitter']:
social_root = f'{dataroot}/{social}'
for dataset in dataset_structure:
print(f'Processing {social}/{dataset}')
current_dir = f'{social_root}/{dataset}'
images = os.listdir(current_dir)
for image in images:
try:
image_key = f'{dataset}/{os.path.splitext(image)[0]}.png'
results_shared[f'{social}/{dataset}/{image}'] = results[image_key]
except KeyError:
image_key = f'{dataset}/{os.path.splitext(image)[0]}.jpg'
results_shared[f'{social}/{dataset}/{image}'] = results[image_key]
results_presocial = {}
for key, value in results.items():
results_presocial[f'PreSocial/{key}'] = value
print(len(results))
print(len(results_presocial))
print(len(results_shared))
print(len(results_shared)+len(results_presocial))
results_all = {**results_presocial, **results_shared}
print(len(results_all))
with open("./classes.pkl", "wb") as f:
pickle.dump(results_all, f, protocol=pickle.HIGHEST_PROTOCOL)
|