| | 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) |
| |
|
