|
import argparse |
|
import csv |
|
import os |
|
|
|
from torchvision.datasets import ImageNet |
|
|
|
|
|
def get_filelist(file_path): |
|
Filelist = [] |
|
for home, dirs, files in os.walk(file_path): |
|
for filename in files: |
|
Filelist.append(os.path.join(home, filename)) |
|
return Filelist |
|
|
|
|
|
def split_by_capital(name): |
|
|
|
new_name = "" |
|
for i in range(len(name)): |
|
if name[i].isupper() and i != 0: |
|
new_name += " " |
|
new_name += name[i] |
|
return new_name |
|
|
|
|
|
def process_imagenet(root, split): |
|
root = os.path.expanduser(root) |
|
data = ImageNet(root, split=split) |
|
samples = [(path, data.classes[label][0]) for path, label in data.samples] |
|
output = f"imagenet_{split}.csv" |
|
|
|
with open(output, "w") as f: |
|
writer = csv.writer(f) |
|
writer.writerows(samples) |
|
|
|
print(f"Saved {len(samples)} samples to {output}.") |
|
|
|
|
|
def process_ucf101(root, split): |
|
root = os.path.expanduser(root) |
|
video_lists = get_filelist(os.path.join(root, split)) |
|
classes = [x.split("/")[-2] for x in video_lists] |
|
classes = [split_by_capital(x) for x in classes] |
|
samples = list(zip(video_lists, classes)) |
|
output = f"ucf101_{split}.csv" |
|
|
|
with open(output, "w") as f: |
|
writer = csv.writer(f) |
|
writer.writerows(samples) |
|
|
|
print(f"Saved {len(samples)} samples to {output}.") |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("dataset", type=str, choices=["imagenet", "ucf101"]) |
|
parser.add_argument("root", type=str) |
|
parser.add_argument("--split", type=str, default="train") |
|
args = parser.parse_args() |
|
|
|
if args.dataset == "imagenet": |
|
process_imagenet(args.root, args.split) |
|
elif args.dataset == "ucf101": |
|
process_ucf101(args.root, args.split) |
|
else: |
|
raise ValueError("Invalid dataset") |
|
|
