Spaces:
Sleeping
Sleeping
| 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) | |