Spaces:
Runtime error
Runtime error
File size: 2,321 Bytes
5ab0373 592e4db 5ab0373 592e4db 5ab0373 592e4db 5ab0373 |
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 |
from dataloader import create_dataloader_aris
from aris import BEAM_WIDTH_DIR
import json
import cv2
import base64
VIDEO_HEIGHT = 700
def load_frames(video_path, preds):
"""Load frames for annotation editing
"""
dataloader, dataset = create_dataloader_aris(video_path, BEAM_WIDTH_DIR, None)
didson = dataset.didson
frames = didson.load_frames(start_frame=0)
frame_info, h, w = get_frame_info(frames, preds)
return frame_info
def get_frame_info(frames, preds):
"""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)
"""
color_map = { fish['id'] : fish['color'] for fish in preds['fish'] }
frame_info = []
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]):
image = cv2.resize(cv2.cvtColor(frame_raw, cv2.COLOR_GRAY2BGR), (w,h))
retval, buffer = cv2.imencode('.jpg', image)
jpg_as_text = base64.b64encode(buffer).decode("utf-8")
frame = {
'annotations': [],
'frame': jpg_as_text
}
for fish in preds['frames'][i]['fish']:
xmin, ymin, xmax, ymax = fish['bbox']
hexx = color_map[fish['fish_id']].lstrip('#')
frame['annotations'].append({
'bbox': {
'left': int(round(xmin * w)),
'right': int(round(xmax * w)),
'top': int(round(ymin * h)),
'bottom': int(round(ymax * h)),
},
'id': str(fish['fish_id']),
'conf': fish['conf']
})
frame_info.append(frame)
return frame_info, h, w |