junction_processing/prepare_junction_maps.py at main · serre-lab/junction_processing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
import os
import sys
import torch
import numpy as np
from PIL import Image
import json
import pandas as pd
import argparse
from src import utils


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=" Junction map for modeling")
    parser.add_argument('config', nargs='?', help="Path to config file")
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--max-workers', type=int, default=None, help='Maximum number of CPU workers')
    args = parser.parse_args()

    # Load config file
    if args.config:
        config_file = args.config if "configs" + os.path.sep in args.config else os.path.join("configs", args.config)
        assert os.path.exists(config_file), f"Cannot find config file: {config_file}"
        config = utils.process_config(config_file)
    else:
        config_file = utils.get_config(sys.argv)
        config = utils.process_config(config_file)

    junction_data = utils.process_junction_data(
        config["junction_data"],
    )

    total_maps = len(junction_data)
    image_names = junction_data.keys()
    total_unique_images = len(image_names)
    print(f"Found {total_unique_images} unique images")

    # Optimize number of workers based on CPU count
    cpu_count = os.cpu_count()
    if args.max_workers:
        config["n_jobs"] = min(args.max_workers, cpu_count)
    else:
        # Leave some cores free for system operations
        config["n_jobs"] = max(1, min(cpu_count - 1, 8))

    # Verify GPU is available
    config["use_gpu_blurring"] = torch.cuda.is_available()
    if config["use_gpu_blurring"]:
        # Print GPU info
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)  # GB
        print(f"Using GPU: {gpu_name} with {gpu_memory:.2f} GB memory")
    else:
        print("GPU not available, exiting.")
        sys.exit(1)

    output_dir = config["assets"]
    image_output_dir = config["example_image_output_dir"]
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(image_output_dir, exist_ok=True)
    os.makedirs(os.path.join(output_dir, config["experiment_name"]), exist_ok=True)

    all_avg_maps, junction_maps = utils.process_junctions(
        junction_data=junction_data,
        config=config,
        image_path=config["image_path"],
        n_jobs=config["n_jobs"]
    )
    # Process visualization for display images if needed
    if config["display_image_keys"]:
        if config["display_image_keys"] == "auto":
            sz_dict = {k: len(v) for k, v in all_avg_maps.items()}
            # arg = np.argsort(list(sz_dict.values()))
            arg = [0, 12, 24, 36, 48, 60, 72, 84, 96, 108]
            config["display_image_keys"] = np.asarray(list(sz_dict.keys()))[arg]

    utils.save_junction_maps(all_avg_maps, junction_maps, output_dir, config["experiment_name"],
                 config["display_image_keys"], config["image_path"], image_output_dir)