CreativeInquiry-DataVis-Project/2_29_24_script.py at main · jschvan/CreativeInquiry-DataVis-Project · 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
113
114
115
116
117
118
119
120
121
122
import cv2
import os
import pandas as pd
import collections
import numpy as np
from queue import Queue


class cell_dat:
    def __init__(self, pos_x, pos_y, velocity_x_8, velocity_y_8):
        self.pos_x = pos_x
        self.pos_y = pos_y
        self.velocity_x_8 = velocity_x_8
        self.velocity_y_8 = velocity_y_8

def csv_scanner(file_name):
    #create pandas dataframe from csv file
    pandas_df = pd.read_csv(file_name)

    #dictionary with frame key and list of cell_dat object values
    frame_cell_dict = collections.defaultdict(list)

    for index, row in pandas_df.iterrows():
        key = row['FRAME']
        template_cell_dat = cell_dat(
            row['pos_x'],
            row['pos_y'],
            row['dt4_n0_dx'],
            row['dt4_n0_dy']
        )
        frame_cell_dict[key].append(template_cell_dat)
    print("scan complete for " + file_name)
    return frame_cell_dict

parent_directory = os.path.join(os.getcwd(), "..")

#setting parameters for video
#tells program where to source its images from
image_folder = 'references/A_01fld07'

#customizes output information for the video
video_name = 'broken_alg_1-7-alt_alg.mp4'
fps = 7
green = [0,255,0]
red = [0,0,255]
blue = [255, 0, 0]


#tells program where to find csv info
#currently using the cnn_dt8 algorithim
raw_csv_file = "csv_files/1-7/spots_velocity.csv"
alg_csv_file = "csv_files/1-7/A_01fld07.csv" #trackmate_dt4 file
interval = 4

#gathers csv data into dictionaries
cell_data_dict = csv_scanner(raw_csv_file)
alg_data_dict = csv_scanner(alg_csv_file)

#collect set of images from given folder and stores quantity of pictures saved
images = [img for img in os.listdir(os.path.join(parent_directory, image_folder)) if img.endswith(".jpg")]
image_qt = len(images)

#saves one individual frame to pull width and height of frame from
frame = cv2.imread(os.path.join(parent_directory + '/' + image_folder, images[0]))
width = frame.shape[1]
height = frame.shape[0]

#initializes video to add images to
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
video = cv2.VideoWriter(os.path.join('../output/output_videos', video_name), fourcc, fps, (width,height))

loop_count = 0
buffer_ct = 3
draw_list = []


for image_num in range(image_qt):

    #generate name for image to be found
    image_name = str(image_num) + ".jpg"

    #open image for drawing points on
    img = cv2.imread( parent_directory + '/' + image_folder + '/' + image_name ,cv2.IMREAD_COLOR)

    #for every object in a given frame
    for item in cell_data_dict[image_num]:
        #parameters: image, center of circle, radius, color, thickness
        cell_coords = (int(item.pos_x), int(item.pos_y))
        cv2.circle(img, cell_coords, 5, green, -1)


    list_of_pairs = []
    for item in alg_data_dict[int(image_num/interval)]:
        #parameters: image, center of circle, radius, color, thickness
        cell_coords = (int(item.pos_x), int(item.pos_y))
        cell_predict = ((int(item.pos_x)+int(item.velocity_x_8)), (int(item.pos_y)+int(item.velocity_y_8)))
        predict_pair = (cell_coords, cell_predict)
        list_of_pairs.append(predict_pair)

    if list_of_pairs not in draw_list:
        if len(draw_list) >= buffer_ct:
            draw_list.pop(0)
        draw_list.append(list_of_pairs)

    for list in draw_list:
        for item in list:
            cv2.line(img, item[0], item[1], red, 5)


    loop_count = loop_count + 1

    cv2.imwrite(os.path.join(parent_directory, 'output/marked/marked_'+ image_name), img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    video.write(cv2.imread(os.path.join('../output/marked', 'marked_'+ image_name)))

cv2.destroyAllWindows()
video.release()

print("Video has been created!")