CreativeInquiry-DataVis-Project/buffer_video_gen.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
123
124
125
126
127
128
129
130
131
import cv2
import os
import csv #can i delete?
import pandas as pd
import collections
import numpy as np


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['dt8_n0_dx'],
            row['dt8_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_01fld01'

#customizes output information for the video
video_name = '8__predict+xy_track.mp4' #can we make this 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 = "references/csv_files/spots_velocity.csv"
alg_csv_file = "references/csv_files/A_01fld01.csv"
interval = 8

#gathers csv data into dictionaries
cell_data_dict = csv_scanner(os.path.join(parent_directory, raw_csv_file))
alg_data_dict = csv_scanner(os.path.join(parent_directory, 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))

#runs through number in range
#does this as opposed to looping through every image in folder
#could this be fixed?

loop_count = 0
buffer_ct = 5
draw_list = []

from queue import Queue
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)

    '''
    for item in cell_data_dict[image_num-loop_count]:
        cell_coords = (int(item.pos_x), int(item.pos_y))
        cell_velocity = (int(item.pos_x) + int(item.velocity_x_8)), (int(item.pos_y) + int(item.velocity_y_8))
        cv2.line(img, cell_coords, cell_velocity, green, 5)
    '''

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