Spaces:
Runtime error
Runtime error
File size: 8,563 Bytes
7a4b92f 592e4db 7a4b92f 592e4db 7a4b92f 592e4db 7a4b92f 1edac8d 7a4b92f 592e4db 7a4b92f |
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 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
import project_path
import json
import cv2
import numpy as np
from tqdm import tqdm
from lib.fish_eye.tracker import Tracker
VERSION = "09/21"
PRED_COLOR = (255, 0, 0) # blue
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BORDER_PAD = 3
LINE_HEIGHT= 22
VIDEO_HEIGHT = 700
INFO_PANE_WIDTH = 100
BOX_THICKNESS = 2
FONT_SCALE = 0.65
FONT_THICKNESS = 1
def generate_video_batches(didson, preds, frame_rate, video_out_path, gp=None, image_meter_width=None, image_meter_height=None, batch_size=1000):
"""Write a visualized video to video_out_path, given a didson object.
"""
if (gp): gp(0, "Generating results video...")
end_frame = didson.info['endframe'] or didson.info['numframes']
out = None # need to wait til we have height and width to instantiate video file
print(preds)
print(preds.keys())
with tqdm(total=end_frame, desc="Generating results video", ncols=0) as pbar:
for i in range(0, end_frame, batch_size):
batch_end = min(end_frame, i+batch_size)
frames = didson.load_frames(start_frame=i, end_frame=batch_end)
vid_frames, h, w = get_video_frames(frames, preds, frame_rate, image_meter_width, image_meter_height, start_frame=i)
if out is None:
out = cv2.VideoWriter(video_out_path, cv2.VideoWriter_fourcc(*'mp4v'), frame_rate, [ int(1.5*w), h ] )
for j, frame in enumerate(vid_frames):
if gp: gp(( (i+j) / end_frame), 'Generating results video...')
out.write(frame)
pbar.update(1)
del frames
del vid_frames
out.release()
def get_video_frames(frames, preds, frame_rate, image_meter_width=None, image_meter_height=None, start_frame=0):
"""Get visualized video frames ready for output, given raw ARIS/DIDSON frames.
Warning: all frames in frames will be stored in memory - careful of OOM errors. Consider processing large files
in batches, such as in generate_video_batches()
Returns:
list(np.ndarray), height (int), width (int)
"""
pred_lengths = { fish['id'] : "%.2fm" % fish['length'] for fish in preds['fish'] }
clip_pr_counts = Tracker.count_dirs(preds)
color_map = { fish['id'] : fish['color'] for fish in preds['fish'] }
# filter JSON, if necessary (for shorter clips)
preds['frames'] = preds['frames'][start_frame:]
vid_frames = []
if len(frames):
# assumes all frames the same size
h, w = frames[0].shape
# enforce a standard size so that text/box thickness is consistent
scale_factor = VIDEO_HEIGHT / h
h = VIDEO_HEIGHT
w = int(scale_factor*w)
num_frames = min(len(frames), len(preds['frames']))
for i, frame_raw in enumerate(frames[:num_frames]):
frame_raw = cv2.resize(cv2.cvtColor(frame_raw, cv2.COLOR_GRAY2BGR), (w,h))
pred = preds['frames'][i]
for fish in pred['fish']:
xmin, ymin, xmax, ymax = fish['bbox']
left = int(round(xmin * w))
right = int(round(xmax * w))
top = int(round(ymin * h))
bottom = int(round(ymax * h))
fish_id = str(fish['fish_id'])
fish_len = pred_lengths[fish['fish_id']]
hexx = color_map[fish['fish_id']].lstrip('#')
color = tuple(int(hexx[i:i+2], 16) for i in (0, 2, 4))
draw_fish(frame_raw, left, right, top, bottom, color, fish_id, fish_len, anno_align="right")
# add axis to frame
frame_raw = add_axis(frame_raw, image_meter_width, image_meter_height)
# add info
frame_info_panel = np.zeros((h, int(0.5*w), 3)).astype(np.uint8)
frame = np.concatenate((frame_info_panel, frame_raw), axis=1)
cv2.putText(frame, f'VERSION: {VERSION}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*4), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
cv2.putText(frame, f'Right count: {clip_pr_counts[0]}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*3), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
cv2.putText(frame, f'Left count: {clip_pr_counts[FONT_THICKNESS]}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*2), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
cv2.putText(frame, f'Other fish: {clip_pr_counts[2]}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*1), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
# cv2.putText(frame, f'Upstream: {preds["upstream_direction"]}', (0, h-1-LINE_HEIGHT*1), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
cv2.putText(frame, f'Frame: {i}', (BORDER_PAD, h-BORDER_PAD-LINE_HEIGHT*0), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
vid_frames.append(frame)
return vid_frames, h, w
def draw_fish(frame, left, right, top, bottom, color, fish_id, fish_len, LINE_HEIGHT=18, anno_align="left"):
cv2.rectangle(frame, (left, top), (right, bottom), color, BOX_THICKNESS)
if anno_align == "left":
anno_align = left
else:
anno_align = right
cv2.putText(frame, fish_id, (anno_align, top), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, color, FONT_THICKNESS, cv2.LINE_AA, False)
cv2.putText(frame, fish_len, (anno_align, bottom+int(LINE_HEIGHT/2)), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE, color, FONT_THICKNESS, cv2.LINE_AA, False)
def add_axis(img, image_meter_width=None, image_meter_height=None):
h, w, c = img.shape
# add black border around image
bordersize_t = 25
bordersize_l = 45
img = cv2.copyMakeBorder(
img,
bottom=bordersize_t,
top=0,
left=bordersize_l,
right=25, # this helps with text getting cut off
borderType=cv2.BORDER_CONSTANT,
value=BLACK
)
# add axis
axis_thickness = 1
img = cv2.line(img, (bordersize_l, h+axis_thickness//2), (w+bordersize_l, h+axis_thickness//2), WHITE, axis_thickness) # x
img = cv2.line(img, (bordersize_l-axis_thickness//2, 0), (bordersize_l-axis_thickness//2, h), WHITE, axis_thickness) # y
# dist between ticks in meters
x_inc = 100
if image_meter_width and image_meter_width > 0:
x_inc = w / image_meter_width / 2 # 0.5m ticks
if image_meter_width > 4:
x_inc *= 2 # 1m ticks
if image_meter_width > 8:
x_inc *= 2 # 2m ticks
# dist between ticks in meters
y_inc = 100
if image_meter_height and image_meter_height > 0:
y_inc = h / image_meter_height / 2 # 0.5m ticks
if image_meter_height > 4:
y_inc *= 2 # 1m ticks
if image_meter_height > 8:
y_inc *= 2 # 2m ticks
if image_meter_height > 12:
y_inc *= 3/2 # 3m ticks
# tick mark labels
def x_label(x):
if image_meter_width and image_meter_width > 0:
if x_inc < w / image_meter_width: # fractional ticks
return "%.1fm" % (x / w * image_meter_width)
return "%.0fm" % (x / w * image_meter_width)
return str(x) # pixels
def y_label(y):
if image_meter_height and image_meter_height > 0:
if y_inc < y / image_meter_height: # fractional ticks
return "%.1fm" % (y / h * image_meter_height)
return "%.0fm" % (y / h * image_meter_height)
return str(y) # pixels
# add ticks
ticksize = 5
x = 0
while x < w:
img = cv2.line(img, (int(bordersize_l+x), h+axis_thickness//2), (int(bordersize_l+x), h+axis_thickness//2+ticksize), WHITE, axis_thickness)
cv2.putText(img, x_label(x), (int(bordersize_l+x), h+axis_thickness//2+LINE_HEIGHT), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE*3/4, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
x += x_inc
y = 0
while y < h:
img = cv2.line(img, (bordersize_l-axis_thickness//2, int(h-y)), (bordersize_l-axis_thickness//2-ticksize, int(h-y)), WHITE, axis_thickness)
ylabel = y_label(y)
txt_offset = 13*len(ylabel)
cv2.putText(img, y_label(y), (bordersize_l-axis_thickness//2-ticksize - txt_offset, int(h-y)), cv2.FONT_HERSHEY_SIMPLEX, FONT_SCALE*3/4, WHITE, FONT_THICKNESS, cv2.LINE_AA, False)
y += y_inc
# resize to original dims
return cv2.resize(img, (w,h)) |