3D-Root-Zonation/getEnhancedSlice.py at main · CIVA-Lab/3D-Root-Zonation · 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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import os
import numpy as np
from PIL import Image
import argparse

def get_enhanced_slice(img_path, out_path, stack_number_given, lim):

    stack_number = ""

    if stack_number_given == "" or stack_number_given =="1":
        stack_number = "-1"
    else:
        stack_number = "-" + stack_number_given

    img_dir = os.path.join(img_path + '/', "TransverseSlice-PS/", 'PS' + stack_number + '/')
    out_dir = os.path.join(out_path + '/', "TransverseSlice-Enh/", 'Enh' + stack_number + '/')

    nn         = 'Enh'             # prefix for output filenames
    lim        = 5                 # same as your MATLAB lim

    # Make sure output directory exists
    os.makedirs(out_dir, exist_ok=True)

    # gather and sort all PNGs
    flist = sorted(f for f in os.listdir(img_dir) if f.lower().endswith('.png'))
    no_frame = len(flist)
    if no_frame == 0:
        raise RuntimeError("No PNG files found in " + img_dir)


    # read first frame to get shape
    first = Image.open(os.path.join(img_dir, flist[0]))
    arr0  = np.array(first, dtype=np.uint8)
    height, width = arr0.shape  # assumes grayscale; if color, use arr0.shape[:2]


    # pre‑allocate the 3D volume
    volume = np.zeros((height, width, no_frame), dtype=np.uint8)

    # load every frame
    for i, fname in enumerate(flist):
        img = Image.open(os.path.join(img_dir, fname))
        volume[:, :, i] = np.array(img, dtype=np.uint8)

    # amount of characters to keep from the end of the original filename
    suffix_length = 9

    # first lim frames: just write them out with the new prefix
    # for i in range(min(lim, no_frame)):
    #     c       = volume[:, :, i]
    #     suffix  = flist[i][-suffix_length:]
    #     outname = nn + suffix
    #     Image.fromarray(c).save(os.path.join(out_dir, outname))
    #     # print("Saved: ", os.path.join(out_dir, outname))


    # remaining frames: compute rem_noise with uint8‐saturating subtract
    for i in range(no_frame - lim):
        # pull out the two frames
        a_uint8 = volume[:, :, i + lim]
        b_uint8 = volume[:, :, i]

        # step 1: cast to signed, subtract, then saturate negatives → diff_sat
        a_int   = a_uint8.astype(np.int16)
        b_int   = b_uint8.astype(np.int16)
        diff_int = a_int - b_int
        diff_sat = np.clip(diff_int, 0, 255)   # MATLAB uint8 underflow → 0

        # step 2: subtract that saturated diff from a_int, then saturate
        rem_int  = a_int - diff_sat
        rem_sat  = np.clip(rem_int, 0, 255).astype(np.uint8)

        # save
        suffix  = flist[i + lim][-suffix_length:]
        outname = nn + suffix
        Image.fromarray(rem_sat).save(os.path.join(out_dir, outname))
        # print("Saved: ", os.path.join(out_dir, outname))


if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--input_path",
        default="./assets/input.mp4",
        help="path to input folder",
    )

    parser.add_argument(
        "--output_path",
        default="./assets/input.mp4",
        help="path to input folder",
    )

    parser.add_argument(
        "--stack_number",
        default="-2/",
        help="number of stack without using predictions",
    )

    parser.add_argument(
        "--limit",
        type=int,
        default=5,
        help="The limit to look to remove the noise in the slice",
    )

    args = parser.parse_args()

    get_enhanced_slice(args.input_path, args.output_path, args.stack_number, lim = args.limit)