| import argparse |
| import json |
| import os |
|
|
| train_rate = 0.6 |
| val_rate = 0.2 |
| test_rate = 0.2 |
|
|
|
|
| def save_json(content, path, filename): |
| if not os.path.exists(path): |
| os.makedirs(path, exist_ok=True) |
| dst_file_name = os.path.join(path, filename) |
| with open(dst_file_name, "w+") as fp: |
| json.dump(content, fp, indent=4, separators=(",", ":")) |
|
|
|
|
| def generate_labels(data_path, output_path): |
| """ |
| Loading a model by name. |
| |
| Args: |
| data_path: path to classification dataset, which must contain `inbody` and `outbody` directories. |
| output_path: path to save labels |
| """ |
|
|
| data_list = [os.path.join(root, x) for root, _, filenames in os.walk(data_path) for x in filenames if "jpg" in x] |
| label_list = [int("outbody" in os.path.basename(os.path.dirname(x))) for x in data_list] |
| data_label_json = [{"image": x, "label": y} for x, y in zip(data_list, label_list)] |
| inbody_list = list(filter(lambda x: x["label"] == 0, data_label_json)) |
| outbody_list = list(filter(lambda x: not (x["label"] == 0), data_label_json)) |
| inbody_train_len = int(len(inbody_list) * train_rate) |
| outbody_train_len = int(len(outbody_list) * train_rate) |
| inbody_val_len = int(len(inbody_list) * (train_rate + val_rate)) |
| outbody_val_len = int(len(outbody_list) * (train_rate + val_rate)) |
| inbody_train_list = inbody_list[:inbody_train_len] |
| outbody_train_list = outbody_list[:outbody_train_len] |
| inbody_val_list = inbody_list[inbody_train_len:inbody_val_len] |
| outbody_val_list = outbody_list[outbody_train_len:outbody_val_len] |
| inbody_test_list = inbody_list[inbody_val_len:] |
| outbody_test_list = outbody_list[outbody_val_len:] |
| train_list = inbody_train_list + outbody_train_list |
| val_list = inbody_val_list + outbody_val_list |
| test_list = inbody_test_list + outbody_test_list |
| save_json(train_list, out_path, "train_samples.json") |
| save_json(val_list, out_path, "val_samples.json") |
| save_json(test_list, out_path, "test_samples.json") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| |
| parser.add_argument( |
| "--datapath", |
| type=str, |
| default=r"/workspace/data/endoscopic_inbody_classification", |
| help="The root path of the inbody classification dataset.", |
| ) |
|
|
| |
| parser.add_argument("--outpath", type=str, default=r"./label", help="The output path of labels.") |
|
|
| args = parser.parse_args() |
| data_path = args.datapath |
| out_path = args.outpath |
|
|
| if not os.path.exists(out_path): |
| os.makedirs(out_path, exist_ok=True) |
| generate_labels(data_path, out_path) |
|
|
