| |
| import os |
| import random |
|
|
| import numpy as np |
|
|
|
|
| def _extract_data_based_dist( |
| data_type: str, image_paths: list, labels: list, dist: list, **params |
| ): |
| def sampling_frames(image_paths: list, labels: list, dist: list, **params): |
| """ |
| Extracting image paths and labels based on given distribution |
| """ |
| total_f = sum(labels) |
| total_r = len(labels) - total_f |
|
|
| assert ( |
| total_f > 0 |
| ), "Number of fake images must be greater than 0 for distribution sampling!" |
| assert ( |
| total_r > 0 |
| ), "Number of real images must be greater than 0 for distribution sampling!" |
|
|
| r_dist, f_dist = dist[0], dist[1] |
|
|
| idxes = sorted(range(0, len(labels)), key=lambda k: labels[k]) |
| r_idxes = idxes[:total_r] |
| f_idxes = idxes[total_r:] |
|
|
| print(f"Original Number of Fake images --- {len(f_idxes)}") |
| print(f"Original Number of Real images --- {len(r_idxes)}") |
|
|
| if int((total_f / f_dist) * r_dist) > total_r: |
| total_f = int((total_r / r_dist) * f_dist) |
| f_idxes = random.sample(f_idxes, total_f) |
| else: |
| total_r = int((total_f / f_dist) * r_dist) |
| r_idxes = random.sample(r_idxes, total_r) |
|
|
| print( |
| f"Number of Fake images --- {len(f_idxes)} given Fake distribution --- {f_dist}" |
| ) |
| print( |
| f"Number of Real images --- {len(r_idxes)} given Real distribution --- {r_dist}" |
| ) |
|
|
| new_idxes = r_idxes + f_idxes |
|
|
| image_paths = [image_paths[i] for i in new_idxes] |
| labels = np.array(labels)[new_idxes] |
|
|
| for k, v in params.items(): |
| if v is not None and len(v): |
| params[k] = [v[i] for i in new_idxes] |
|
|
| return image_paths, labels, params |
|
|
| def sampling_videos(image_paths: list, labels: list, dist: list, **params): |
| """ |
| Extracting image paths and labels based on given distribution for video data |
| """ |
| f_vid_ids = [] |
| r_vid_ids = [] |
| for ip in image_paths: |
| faketype = ip.split("/")[8] |
| vid_id = os.path.dirname(ip) |
| if ( |
| faketype == "real_videos" |
| or "real" in faketype |
| or "original" in faketype |
| ): |
| r_vid_ids.append(vid_id) |
| else: |
| f_vid_ids.append(vid_id) |
|
|
| f_vid_ids = list(set(f_vid_ids)) |
| r_vid_ids = list(set(r_vid_ids)) |
| total_f = len(f_vid_ids) |
| total_r = len(r_vid_ids) |
|
|
| assert ( |
| total_f > 0 |
| ), "Number of fake videos must be greater than 0 for distribution sampling!" |
| assert ( |
| total_r > 0 |
| ), "Number of real videos must be greater than 0 for distribution sampling!" |
| print(f"Original Number of Fake videos --- {total_f}") |
| print(f"Original Number of Real videos --- {total_r}") |
|
|
| r_dist, f_dist = dist[0], dist[1] |
|
|
| if int((total_f / f_dist) * r_dist) > total_r: |
| total_f = int((total_r / r_dist) * f_dist) |
| f_vid_ids = random.sample(f_vid_ids, total_f) |
| else: |
| total_r = int((total_f / f_dist) * r_dist) |
| r_vid_ids = random.sample(r_vid_ids, total_r) |
|
|
| print( |
| f"Number of Fake videos --- {len(f_vid_ids)} given Fake distribution --- {f_dist}" |
| ) |
| print( |
| f"Number of Real videos --- {len(r_vid_ids)} given Real distribution --- {r_dist}" |
| ) |
|
|
| vid_ids = r_vid_ids + f_vid_ids |
| new_idxes = [] |
|
|
| for i in range(len(labels)): |
| ip = image_paths[i] |
| vid_id = "/".join([ip.split("/")[-3], ip.split("/")[-2]]) |
| if vid_id in vid_ids: |
| new_idxes.append(i) |
|
|
| image_paths = [image_paths[i] for i in new_idxes] |
| labels = np.array(labels)[new_idxes] |
|
|
| for k, v in params.items(): |
| if v is not None and len(v): |
| params[k] = [v[i] for i in new_idxes] |
|
|
| return image_paths, labels, params |
|
|
| if data_type == "image": |
| return sampling_frames(image_paths, labels, dist, **params) |
| else: |
| return sampling_videos(image_paths, labels, dist, **params) |
|
|
