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