Spaces:
Runtime error
Runtime error
| import json | |
| import cv2 | |
| import base64 | |
| VIDEO_HEIGHT = 700 | |
| def init_frames(dataset, preds, index, gp=None): | |
| """Load frames for annotation editing | |
| Returns: | |
| list({ | |
| frame: frame image as base64 string, | |
| annotations: list( | |
| bbox: dict of int defining bounding box {left, right, top, bottom}, | |
| id: id of fish as int, | |
| conf: confidence in bbox as float | |
| ) | |
| }) | |
| """ | |
| images = dataset.didson.load_frames(start_frame=0, end_frame=1) | |
| # assumes all frames the same size | |
| h, w = images[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) | |
| annotations = [] | |
| if gp: gp(0, "Extracting Frames") | |
| if len(preds['frames']): | |
| end_index = min(index+1000, len(preds['frames'])) | |
| for i, frame_info in enumerate(preds['frames'][index:end_index]): | |
| if gp: gp((index + i)/len(preds['frames']), "Extracting Frames") | |
| # Extract frames | |
| img_raw = dataset.didson.load_frames(start_frame=index+i, end_frame=index+i+1)[0] | |
| image = cv2.resize(cv2.cvtColor(img_raw, cv2.COLOR_GRAY2BGR), (w, h)) | |
| #cv2.imwrite("annotation_frame_dir/" + str(i) + ".jpg", image) | |
| retval, buffer = cv2.imencode('.jpg', image) | |
| b64 = base64.b64encode(buffer).decode("utf-8") | |
| # Extract annotations | |
| frame = { | |
| 'annotations': [], | |
| 'base64': b64 | |
| } | |
| for fish in frame_info['fish']: | |
| xmin, ymin, xmax, ymax = fish['bbox'] | |
| 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'] | |
| }) | |
| annotations.append(frame) | |
| return annotations, end_index | |