Spaces:
Runtime error
Runtime error
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_path): | |
"""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_path) | |
return frame_info | |
def get_frame_info(frames, preds_path): | |
"""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) | |
""" | |
preds = json.load(open(preds_path, 'r')) | |
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 |