import project_path import numpy as np import cv2 import os from collections import namedtuple, defaultdict import struct from PIL import Image from tqdm import tqdm import datetime from decimal import Decimal, ROUND_HALF_UP import json import pytz from copy import deepcopy from multiprocessing import Pool import math import lib.fish_eye.pyARIS as pyARIS from lib.fish_eye.tracker import Tracker BEAM_WIDTH_DIR = 'lib/fish_eye/beam_widths/' ImageData = namedtuple('ImageData', [ 'pixel_meter_size', 'xdim', 'ydim', 'x_meter_start', 'y_meter_start', 'x_meter_stop', 'y_meter_stop', 'sample_read_rows', 'sample_read_cols', 'image_write_rows', 'image_write_cols' ]) def FastARISRead(ARIS_data, start_frame, end_frame): """ Just read in the ARIS frame, and not the other meta data. """ FrameSize = ARIS_data.SamplesPerChannel*ARIS_data.NumRawBeams frames = np.empty([end_frame-start_frame, ARIS_data.SamplesPerChannel, ARIS_data.NumRawBeams], dtype=np.uint8) with open(ARIS_data.filename, 'rb') as data: for i, j in enumerate(range(start_frame, end_frame)): data.seek(j*(1024+(FrameSize))+2048, 0) raw_data = struct.unpack("%dB" % FrameSize, data.read(FrameSize)) frames[i] = np.fliplr(np.reshape( raw_data, [ARIS_data.SamplesPerChannel, ARIS_data.NumRawBeams])) # Close the data file data.close() return frames def get_info(aris_fp, beam_width_dir=BEAM_WIDTH_DIR): """ Return: image_meter_width, image_meter_height, fps """ ARISdata, aris_frame = pyARIS.DataImport(aris_fp) beam_width_data = pyARIS.load_beam_width_data(aris_frame, beam_width_dir=beam_width_dir)[0] min_pixel_size = pyARIS.get_minimum_pixel_meter_size(aris_frame, beam_width_data) sample_length = aris_frame.sampleperiod * 0.000001 * aris_frame.soundspeed / 2 pixel_meter_size = max(min_pixel_size, sample_length) xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop = pyARIS.compute_image_bounds( pixel_meter_size, aris_frame, beam_width_data, additional_pixel_padding_x=0, additional_pixel_padding_y=0 ) return pixel_meter_size * xdim, pixel_meter_size * ydim, aris_frame.framerate def write_frames(aris_fp, out_dir, cb=None, max_mb=-1, beam_width_dir=BEAM_WIDTH_DIR, bg_out_dir=None, num_workers=0): """ Write all frames from an ARIS file to disk, using our 3-channel format: (raw img, blurred & mean subtracted img, optical flow approximation) Args: aris_fp: path to aris file out_dir: directory for frame extraction. frames will be named 0.jpg, 1.jpg, ... {n}.jpg cb: a callback function for updating progress max_mb: maximum amount of the file to be processed, in megabytes beam_width_dir: location of ARIS camera information bg_out_dir: where to write the background frame; None disables writing Return: (float) image_meter_width - the width of each image, in meters (float) image_meter_height (float) fps """ # Load in the ARIS file ARISdata, aris_frame = pyARIS.DataImport(aris_fp) if cb: cb(2, msg="Decoding ARIS data...") beam_width_data = pyARIS.load_beam_width_data(aris_frame, beam_width_dir=beam_width_dir)[0] # What is the meter resolution of the smallest sample? min_pixel_size = pyARIS.get_minimum_pixel_meter_size(aris_frame, beam_width_data) # What is the meter resolution of the sample length? sample_length = aris_frame.sampleperiod * 0.000001 * aris_frame.soundspeed / 2 # Choose the size of a pixel (or hard code it to some specific value) pixel_meter_size = max(min_pixel_size, sample_length) # Determine the image dimensions xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop = pyARIS.compute_image_bounds( pixel_meter_size, aris_frame, beam_width_data, additional_pixel_padding_x=0, additional_pixel_padding_y=0 ) # Compute the mapping from the samples to the image sample_read_rows, sample_read_cols, image_write_rows, image_write_cols = pyARIS.compute_mapping_from_sample_to_image( pixel_meter_size, xdim, ydim, x_meter_start, y_meter_start, aris_frame, beam_width_data ) image_data = ImageData( pixel_meter_size, xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop, sample_read_rows, sample_read_cols, image_write_rows, image_write_cols ) start_frame = 0 end_frame = ARISdata.FrameCount bytes_per_frame = 1024+ARISdata.SamplesPerChannel*ARISdata.NumRawBeams print("ARIS bytes per frame", bytes_per_frame) img_bytes_per_frame = image_data.ydim * image_data.xdim * 4 # for fp32 frames print("Image bytes per frame", img_bytes_per_frame) max_bytes = max(bytes_per_frame, img_bytes_per_frame) if max_mb > 0: max_frames = int(max_mb*1000000 / (max_bytes)) if end_frame > max_frames: end_frame = max_frames # use a max of 4gb per batch to avoid memory errors (16gb RAM on a g4dn.xlarge) batch_size = 1000 # int(4000*1000000 / (max_bytes)) clips = [[pos, pos+batch_size+1] for pos in range(0, end_frame, batch_size)] clips[-1][1] = ARISdata.FrameCount print("Batch size:", batch_size) with tqdm(total=(end_frame-start_frame-1), desc="Extracting frames", ncols=0) as pbar: # compute info for bg subtraction using first batch # TODO: make this a sliding window mean_blurred_frame, mean_normalization_value = write_frame_range(ARISdata, image_data, out_dir, clips[0][0], clips[0][1], None, None, cb, pbar) # do rest of batches in parallel if num_workers > 0: args = [ (ARISdata, image_data, out_dir, start, end, mean_blurred_frame, mean_normalization_value, cb) for (start, end) in clips[1:] ] # TODO: can't pass pbar to thread with Pool(num_workers) as pool: results = [ pool.apply_async(write_frame_range, arg) for arg in args ] results = [ r.get() for r in results ] # need this call to block on thread execution pbar.update(sum([ arg[4] - arg[3] for arg in args ])) else: for j, (start, end) in enumerate(clips[1:]): write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_frame, mean_normalization_value, cb, pbar) if bg_out_dir is not None: bg_img = (mean_blurred_frame * 255).astype(np.uint8) out_fp = os.path.join(bg_out_dir, 'bg_start.jpg') Image.fromarray(bg_img).save(out_fp, quality=95) return pixel_meter_size * xdim, pixel_meter_size * ydim, aris_frame.framerate def write_frame_range(ARISdata, image_data, out_dir, start, end, mean_blurred_frame=None, mean_normalization_value=None, cb=None, pbar=None): try: frames = np.zeros([end-start, image_data.ydim, image_data.xdim], dtype=np.uint8) frames[:, image_data.image_write_rows, image_data.image_write_cols] = FastARISRead(ARISdata, start, end)[:, image_data.sample_read_rows, image_data.sample_read_cols] except: print("Error extracting frames from", ARISdata.filename, "during batch", i) return blurred_frames = frames.astype(np.float32) for i in range(frames.shape[0]): blurred_frames[i] = cv2.GaussianBlur( blurred_frames[i], (5,5), 0 ) if mean_blurred_frame is None: mean_blurred_frame = blurred_frames.mean(axis=0) blurred_frames -= mean_blurred_frame if mean_normalization_value is None: mean_normalization_value = np.max(np.abs(blurred_frames)) blurred_frames /= mean_normalization_value blurred_frames += 1 blurred_frames /= 2 # Because of the optical flow computation, we only go to end_frame - 1 for i, frame_offset in enumerate(range(start, end - 1)): frame_image = np.dstack([ frames[i] / 255, blurred_frames[i], np.abs(blurred_frames[i+1] - blurred_frames[i]) ]).astype(np.float32) frame_image = (frame_image * 255).astype(np.uint8) out_fp = os.path.join(out_dir, f'{start+i}.jpg') # = frame_offset.jpg? Image.fromarray(frame_image).save(out_fp, quality=95) if pbar: pbar.update(1) if cb: pct = 2 + int( (start+i) / (end_frame - start_frame - 1) * 98) cb(pct, msg=pbar.__str__()) return mean_blurred_frame, mean_normalization_value def prep_for_mm(json_data): """Prepare json results for writing to a manual marking file.""" json_data = deepcopy(json_data) # map fish id -> [ (bbox, frame_num), (bbox, frame_num), ... ] tracks = defaultdict(list) for frame in json_data['frames']: for bbox in frame['fish']: tracks[bbox['fish_id']].append((bbox['bbox'], frame['frame_num'])) # find frame number for manual marking # look for first time a track crosses the center # if it never crosses the center, use the closest box to the center mm_frame_nums = {} for f_id, track in tracks.items(): # keep track of frame closest to the center closest_frame = 0 closest_dist = 1.0 for i, (box, frame) in enumerate(track): x = (box[0] + box[2]) / 2.0 if i > 0: last_x = (track[i-1][0][0] + track[i-1][0][2]) / 2.0 if (x < 0.5 and last_x >= 0.5) or (last_x < 0.5 and x >= 0.5): closest_frame = frame break dist = abs(x - 0.5) if dist < closest_dist: closest_frame = frame closest_dist = dist mm_frame_nums[f_id] = closest_frame # sort tracks by their frame numbers and re-key # IDs are 1-indexed id_frame = [ (k, v) for k,v in mm_frame_nums.items() ] id_frame = sorted(id_frame, key=lambda x: x[1]) id_map = {} for i, (f_id, frame) in enumerate(id_frame, start=1): id_map[f_id] = i # map IDs and keep frame['fish'] sorted by ID for i, frame in enumerate(json_data['frames']): new_frame_entries = [] for frame_entry in frame['fish']: frame_entry['fish_id'] = id_map[frame_entry['fish_id']] new_frame_entries.append(frame_entry) frame['fish'] = sorted(new_frame_entries, key=lambda k: k['fish_id']) # store manual marking frame and re-map 'fish' field for fish in json_data['fish']: fish['marking_frame'] = mm_frame_nums[fish['id']] # mm_frame_nums refers to old IDs fish['id'] = id_map[fish['id']] json_data['fish'] = sorted(json_data['fish'], key=lambda x: x['id']) return json_data def add_metadata_to_result(aris_fp, json_data, beam_width_dir=BEAM_WIDTH_DIR): """ Return: dictionary, for manual marking """ metadata = {} metadata["FILE_NAME"] = aris_fp ARISdata, frame = pyARIS.DataImport(aris_fp) metadata["FRAME_RATE"] = frame.framerate # Load in the beam width information beam_width_data, camera_type = pyARIS.load_beam_width_data(frame, beam_width_dir=beam_width_dir) # What is the meter resolution of the smallest sample? min_pixel_size = pyARIS.get_minimum_pixel_meter_size(frame, beam_width_data) # What is the meter resolution of the sample length? sample_length = frame.sampleperiod * 0.000001 * frame.soundspeed / 2 # Choose the size of a pixel pixel_meter_size = max(min_pixel_size, sample_length) # Determine the image dimensions xdim, ydim, x_meter_start, y_meter_start, x_meter_stop, y_meter_stop = pyARIS.compute_image_bounds( pixel_meter_size, frame, beam_width_data, additional_pixel_padding_x=0, additional_pixel_padding_y=0 ) # Compute the mapping from the samples to the image sample_read_rows, sample_read_cols, image_write_rows, image_write_cols = pyARIS.compute_mapping_from_sample_to_image( pixel_meter_size, xdim, ydim, x_meter_start, y_meter_start, frame, beam_width_data ) marking_mapping = dict(zip(zip(image_write_rows, image_write_cols), zip(sample_read_rows, sample_read_cols))) # Manual marking format rounds 0.5 to 1 instead of 0 in IEEE 754 def round(number, ndigits=0): return float(Decimal(number).quantize(ndigits, ROUND_HALF_UP)) right, left, none = Tracker.count_dirs(json_data) metadata["UPSTREAM_FISH"] = left # TODO metadata["DOWNSTREAM_FISH"] = right # TODO metadata["NONDIRECTIONAL_FISH"] = none # TODO metadata["TOTAL_FISH"] = metadata["UPSTREAM_FISH"] + metadata["DOWNSTREAM_FISH"] + metadata["NONDIRECTIONAL_FISH"] metadata["TOTAL_FRAMES"] = ARISdata.FrameCount metadata["EXPECTED_FRAMES"] = -1 # What is this? metadata["TOTAL_TIME"] = str(datetime.timedelta(seconds=round(metadata["TOTAL_FRAMES"]/metadata["FRAME_RATE"]))) metadata["EXPECTED_TIME"] = str(datetime.timedelta(seconds=round(metadata["EXPECTED_FRAMES"]/metadata["FRAME_RATE"]))) metadata["UPSTREAM_MOTION"] = 'Right To Left' or 'Left To Right' #TODO metadata["COUNT_FILE_NAME"] = 'N/A' metadata["EDITOR_ID"] = 'N/A' metadata["INTENSITY"] = f'{round(frame.intensity, 1):.1f} dB' # Missing metadata["THRESHOLD"] = f'{round(frame.threshold, 1):.1f} dB' # Missing metadata["WINDOW_START"] = round(frame.windowstart, 2) metadata["WINDOW_END"] = round(frame.windowstart + frame.windowlength, 2) metadata["WATER_TEMP"] = f'{int(round(frame.watertemp))} degC' s = f'''''' upstream_motion_map = {} if (metadata["UPSTREAM_MOTION"] == 'Left To Right'): upstream_motion_map = { 'right': ' Up', 'left': 'Down', 'none': ' N/A', } elif (metadata["UPSTREAM_MOTION"] == 'Right To Left'): upstream_motion_map = { 'left': ' Up', 'right': 'Down', 'none': ' N/A', } def get_entry(fish): if 'marking_frame' in fish: frame_num = fish['marking_frame'] entry = None for json_frame in json_data['frames']: if json_frame['frame_num'] == frame_num: for json_frame_entry in json_frame['fish']: if json_frame_entry['fish_id'] == fish['id']: json_frame_entry = json_frame_entry.copy() json_frame_entry['frame_num'] = frame_num return json_frame_entry else: print("Warning: JSON not correctly formatted for manual marking creation. Use aris.prep_for_mm()") entries = [] for json_frame in json_data['frames']: for json_frame_entry in json_frame['fish']: if json_frame_entry['fish_id'] == fish['id']: entries.append({'frame_num': json_frame['frame_num'], **json_frame_entry}) entry = entries[len(entries)//2] return entry print("Error, could not find entry for", fish) return None # TODO better error handling entries = [] for fish in json_data['fish']: entry = get_entry(fish) entry['length'] = fish['length']*100 entry['direction'] = fish['direction'] entry['travel_dist'] = fish['travel_dist'] entry['start_frame_index'] = fish['start_frame_index'] entry['end_frame_index'] = fish['end_frame_index'] entries.append(entry) metadata["FISH"] = [] for entry in sorted(entries, key=lambda x: x['fish_id']): frame_num = entry['frame_num'] frame = pyARIS.FrameRead(ARISdata, frame_num) y = (entry['bbox'][1]+entry['bbox'][3])/2 x = (entry['bbox'][0]+entry['bbox'][2])/2 h = np.max(image_write_rows) w = np.max(image_write_cols) # TODO actually fix this try: bin_num, beam_num = marking_mapping[(round(y*h), round(x*w))] except: bin_num = 0 beam_num = 0 fish_entry = {} fish_entry['FILE'] = 1 fish_entry['TOTAL'] = entry['fish_id'] fish_entry['FRAME_NUM'] = entry['frame_num'] fish_entry['START_FRAME'] = entry['start_frame_index'] fish_entry['END_FRAME'] = entry['end_frame_index'] fish_entry['NBR_FRAMES'] = entry['end_frame_index'] + 1 - entry['start_frame_index'] fish_entry['TRAVEL'] = entry['travel_dist'] fish_entry['DIR'] = upstream_motion_map[entry['direction']] fish_entry['R'] = bin_num * pixel_meter_size + frame.windowstart fish_entry['THETA'] = beam_width_data['beam_center'][beam_num] fish_entry['L'] = entry['length'] fish_entry['DR'] = -1.0 # What is this? fish_entry['LDR'] = -1.0 # What is this? fish_entry['ASPECT'] = -1.0 # What is this? TIME, DATE = datetime.datetime.fromtimestamp(frame.sonartimestamp/1000000, pytz.timezone('UTC')).strftime('%H:%M:%S %Y-%m-%d').split() fish_entry['TIME'] = TIME fish_entry['DATE'] = DATE fish_entry['LATITUDE'] = frame.latitude or 'N 00 d 0.00000 m' fish_entry['LONGITUDE'] = frame.longitude or 'E 000 d 0.00000 m' fish_entry['PAN'] = frame.sonarpan if math.isnan(fish_entry['PAN']): fish_entry['PAN'] = "nan" fish_entry['TILT'] = frame.sonartilt if math.isnan(fish_entry['TILT']): fish_entry['TILT'] = "nan" fish_entry['ROLL'] = frame.roll # May be wrong number but sonarroll was NaN fish_entry['SPECIES'] = 'Unknown' fish_entry['MOTION'] = 'Running <-->' fish_entry['Q'] = -1 #5 # I don't know what this is or where it comes from fish_entry['N'] = -1 #1 # I don't know what this is or where it comes from fish_entry['COMMENT'] = '' metadata["FISH"].append(fish_entry) # What are these? # Maybe the date and time range for the recording? first_frame = pyARIS.FrameRead(ARISdata, 0) last_frame = pyARIS.FrameRead(ARISdata, metadata["TOTAL_FRAMES"]-1) start_time, start_date = datetime.datetime.fromtimestamp(first_frame.sonartimestamp/1000000, pytz.timezone('UTC')).strftime('%H:%M:%S %Y-%m-%d').split() end_time, end_date = datetime.datetime.fromtimestamp(last_frame.sonartimestamp/1000000, pytz.timezone('UTC')).strftime('%H:%M:%S %Y-%m-%d').split() metadata["DATE"] = start_date metadata["START"] = start_time metadata["END"] = end_time json_data['metadata'] = metadata return json_data def create_metadata_table(result, table_headers, info_headers): if 'metadata' in result: metadata = result['metadata'] else: metadata = { 'FISH': [] } # Calculate detection dropout for fish in metadata['FISH']: count = 0 for frame in result['frames'][fish['START_FRAME']:fish['END_FRAME']+1]: for ann in frame['fish']: if ann['fish_id'] == fish['TOTAL']: count += 1 fish['DETECTION_DROPOUT'] = 1 - count / (fish['END_FRAME'] + 1 - fish['START_FRAME']) # Create fish table table = [] for fish in metadata["FISH"]: row = [] for header in table_headers: row.append(fish[header]) table.append(row) if len(metadata["FISH"]) == 0: row = [] for header in table_headers: row.append("-") table.append(row) # Create info table info = [] for field in info_headers: field_name = "**" + field + "**" if field in metadata: info.append([field_name, str(metadata[field])]) else: info.append([field_name, ""]) if 'hyperparameters' in metadata: for param_name in metadata['hyperparameters']: info.append(['**' + param_name + '**', str(metadata['hyperparameters'][param_name])]) return table, info def create_manual_marking(results, out_path=None): """ Return: string, full contents of manual marking """ metadata = results['metadata'] s = f''' Total Fish = {metadata["TOTAL_FISH"]} Upstream = {metadata["UPSTREAM_FISH"]} Downstream = {metadata["DOWNSTREAM_FISH"]} ?? = {metadata["NONDIRECTIONAL_FISH"]} Total Frames = {metadata["TOTAL_FRAMES"]} Expected Frames = {metadata["EXPECTED_FRAMES"]} Total Time = {metadata["TOTAL_TIME"]} Expected Time = {metadata["EXPECTED_TIME"]} Upstream Motion = {metadata["UPSTREAM_MOTION"]} Count File Name: {metadata["COUNT_FILE_NAME"]} Editor ID = {metadata["EDITOR_ID"]} Intensity = {metadata["INTENSITY"]} Threshold = {metadata["THRESHOLD"]} Window Start = {metadata["WINDOW_START"]:.2f} Window End = {metadata["WINDOW_END"]:.2f} Water Temperature = {metadata["WATER_TEMP"]} *** Manual Marking (Manual Sizing: Q = Quality, N = Repeat Count) *** File Total Frame# Dir R (m) Theta L(cm) dR(cm) L/dR Aspect Time Date Latitude Longitude Pan Tilt Roll Species Motion Q N Comment --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- ''' for fish in metadata["FISH"]: entry = {} for field in fish.keys(): if fish[field] == "nan": entry[field] = math.nan else: entry[field] = fish[field] s += f'{entry["FILE"]:>4} {entry["TOTAL"]:>5} {entry["FRAME_NUM"]:>6} {entry["DIR"]:>3} {entry["R"]:>6.2f} {entry["THETA"]:>6.1f} {entry["L"]:>6.1f} {entry["DR"]:>6.1f} {entry["LDR"]:>6.2f} {entry["ASPECT"]:>6.1f} {entry["TIME"]:>8} {entry["DATE"]:>10} {entry["LATITUDE"]:>17} {entry["LONGITUDE"]:>18} {entry["PAN"]:>7.2f} {entry["TILT"]:>7.2f} {entry["ROLL"]:>7.2f} {entry["SPECIES"]:>8} {entry["MOTION"]:>37} {entry["Q"]:>5} {entry["N"]:>2} {entry["COMMENT"]}\n' s += f''' *** Source File Key *** 1. Source File Name: {metadata["FILE_NAME"]} Source File Date: {metadata["DATE"]} Source File Start: {metadata["START"]} Source File End: {metadata["END"]} Settings Upstream: {metadata["UPSTREAM_MOTION"]} Default Mark Direction: Upstream Editor ID: {metadata["EDITOR_ID"]} Show Marks: ?? Show marks for ?? seconds Loop for ?? seconds ''' if out_path: with open(out_path, 'w') as f: f.write(s) return s