File size: 5,541 Bytes
2689208
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8df7617
2689208
 
 
 
 
 
 
072e991
2689208
 
 
 
 
1f21e52
 
 
 
 
 
 
 
 
2689208
 
 
 
 
 
 
8e134e8
 
2689208
8e134e8
 
2689208
 
8e134e8
2689208
8df7617
 
 
 
 
072e991
 
2689208
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b568bd7
2689208
 
b568bd7
 
 
2689208
 
 
 
b568bd7
 
2689208
 
 
 
072e991
2689208
 
072e991
2689208
 
 
 
 
 
 
072e991
 
 
 
2689208
 
072e991
 
2689208
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
072e991
2689208
 
 
 
 
 
 
 
 
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
import pandas as pd
from collections import defaultdict
from dotenv import load_dotenv
import os
from PIL import Image, ImageDraw
import math
import json
import random

class StackEntry:
    def __init__(self):
        self.images = []
        self.objects = []
    def add_image(self, image):
        self.images.append(image)
    def add_object(self, object):
        self.objects.append(object)
    def sort(self):
        self.images.sort(key=lambda x: x.focus_value)

def get_neighbours(img, x, y, dimensions):
    neighbour_candidates = [(-1,-1), (0, -1), (1, -1), (-1, 0), (1,0), (-1,1), (0,1), (1,1)]

    width, height = img.size

    neighbours = []
    for x_offset, y_offset in neighbour_candidates:
        neighbour_x = x + x_offset * dimensions
        neighbour_y = y + y_offset * dimensions

        if neighbour_x >= 0 and neighbour_x + dimensions <= width and      neighbour_y >= 0 and neighbour_y + dimensions <= height:
            box = [neighbour_x, neighbour_y, neighbour_x + dimensions, neighbour_y + dimensions]
            neighbours.append((neighbour_x, neighbour_y, img.crop(box)))
        else:
            neighbours.append(None)
    return neighbours

def extract_object_tiles(obj, stack_images, in_folder, threshold = 0.25):
    x_start = int(obj.x_min / size) * size
    x_end = int(math.ceil(obj.x_max / size)) * size
    y_start = int(obj.y_min / size) * size
    y_end = int(math.ceil(obj.y_max / size)) * size

    tiles = []

    focus_stack_images = list(map(lambda x: (x, Image.open(os.path.join(in_folder, x.file_path))), stack_images))

    # Get tiles of the image that contain bounding box of object
    for y in range(y_start, y_end, size):
        for x in range(x_start, x_end, size):
            
            if compute_overlap([x, y, x + size, y + size], [obj.x_min, obj.y_min, obj.x_max, obj.y_max]) > size * size * threshold:
                stack = []
                for row, img in focus_stack_images:
                    box = [x, y, x + size, y + size]
                    crop = img.crop(box)
                
                    neighbours = get_neighbours(img, x, y, size)
                    stack.append((row, box[:2], crop, neighbours))
                tiles.append(stack)
    return tiles


def save_tile(original_file_path, out_dir, x : int, y : int, img, overwrite = False):
    path, file_name = os.path.split(original_file_path)
    name, ext = os.path.splitext(file_name)

    out_path = os.path.join(out_dir, path)
    save_to = os.path.join(out_path, f'{name}_{x}_{y}{ext}')

    if not os.path.exists(out_path):
        os.makedirs(out_path)
    if overwrite or not os.path.exists(save_to):
        img.save(save_to)
    return os.path.join(path, f'{name}_{x}_{y}{ext}')

def compute_overlap(rect1, rect2):
    dx = min(rect1[2], rect2[2]) - max(rect1[0], rect2[0])
    dy = min(rect1[3], rect2[3]) - max(rect1[1], rect2[1])
    return dx * dy

def save_obj_tiles(obj, out_folder, in_folder, stack_images):
    extracted = extract_object_tiles(obj, stack_images, in_folder)
    z_stacks = []
    for z_stack in extracted:
        z_stack_images = []
        for row, box, img, neigbours in z_stack:

            neighbours = []

            image_path = save_tile(row.file_path, out_folder, box[0], box[1], img)
            for neighbour in neigbours:
                n_path = None
                if neighbour:
                    x, y, n_img = neighbour
                    n_path = save_tile(row.file_path, out_folder, x, y, n_img)
                neighbours.append(n_path)

            z_stack_images.append({
                "focus_value": row["focus_value"],
                "image_path": image_path,
                "neighbours": neighbours,
                "original_filename": row["file_name"],
                "scan_uuid": row["uuid"],
                "study_id": row["study_id"],
            })
        z_stacks.append({ 
            "best_index": None,
            "images" : z_stack_images,
            "obj_name": obj["name"],
            "stack_id": obj["stack_id"],
        })

    return z_stacks

def save_stack(stack, out_folder, in_folder):
    z_stacks = []
    for obj in stack.objects:
        z_stacks.extend(save_obj_tiles(obj, out_folder, in_folder, stack.images))
    return z_stacks


if __name__ == "__main__":
    load_dotenv()
    print("Geting environment variables...")
    size = int(os.getenv('IMG_SIZE'))
    root_in = os.getenv('ROOT_IN')

    print(f'img_size: ')
    print(f'in_folder: {root_in}')

    print("Loading data from csv files...")
    objects = pd.read_csv("out/test_objects.csv", index_col=0)
    stacks = pd.read_csv("out/test_stacks.csv", index_col=0)


    stacks_dict = defaultdict(lambda: StackEntry())
    
    print("Building internal datastructure...")
    # adding images to dict
    for (index, row) in stacks.iterrows():
        stacks_dict[row.stack_id].add_image(row)

    for values in stacks_dict.values():
        values.sort()

    # adding objects
    for (index, row) in objects.iterrows():
        stacks_dict[row.stack_id].add_object(row)

    out_folder = "out"
    z_stacks = []

    print("Generating image tiles and writing them to file...")
    for stack in stacks_dict.values():
        z_stacks.extend(save_stack(stack,"out", root_in))

    # randomize z_stacks
    print("Shuffling data...")
    random.shuffle(z_stacks)

    print("Writing meta-data for annotation to file...")
    with open(os.path.join(out_folder, "data.json"), 'w') as file:
        file.write(json.dumps(z_stacks))