Spaces:
Runtime error
Runtime error
| import argparse | |
| import torch | |
| import os | |
| import json | |
| from tqdm import tqdm | |
| import project_subpath | |
| from backend.InferenceConfig import InferenceConfig | |
| from backend.dataloader import create_dataloader_frames_only | |
| from backend.inference import do_full_tracking, setup_model, do_detection | |
| def main(args, config=InferenceConfig(), verbose=True): | |
| """ | |
| Perform inference on a directory of frames and saves the tracking json result | |
| Args: | |
| frames (str): Path to frame directory. Required. | |
| metadata (str): Path to metadata directory. Required. | |
| output (str): Path to output directory. Required. | |
| weights (str): Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt | |
| """ | |
| print("In task...") | |
| print("Cuda available in task?", torch.cuda.is_available()) | |
| print("Config:", config.to_dict()) | |
| dirname = args.frames | |
| loc = args.location | |
| in_loc_dir = os.path.join(dirname, loc) | |
| out_dir = os.path.join(args.output, loc, "tracker", "data") | |
| metadata_path = os.path.join(args.metadata, loc + ".json") | |
| os.makedirs(out_dir, exist_ok=True) | |
| print(in_loc_dir) | |
| print(out_dir) | |
| print(metadata_path) | |
| # run detection + tracking | |
| model, device = setup_model(args.weights) | |
| seq_list = os.listdir(in_loc_dir) | |
| idx = 1 | |
| with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar: | |
| for seq in seq_list: | |
| pbar.update(1) | |
| pbar.set_description("Processing " + seq) | |
| if verbose: | |
| print(" ") | |
| print("(" + str(idx) + "/" + str(len(seq_list)) + ") " + seq) | |
| print(" ") | |
| idx += 1 | |
| in_seq_dir = os.path.join(in_loc_dir, seq) | |
| infer_seq(in_seq_dir, out_dir, config, seq, model, device, metadata_path, verbose) | |
| def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path, verbose): | |
| #progress_log = lambda p, m: 0 | |
| image_meter_width = -1 | |
| image_meter_height = -1 | |
| with open(metadata_path, 'r') as f: | |
| json_object = json.loads(f.read()) | |
| for seq in json_object: | |
| if seq['clip_name'] == seq_name: | |
| image_meter_width = seq['x_meter_stop'] - seq['x_meter_start'] | |
| image_meter_height = seq['y_meter_stop'] - seq['y_meter_start'] | |
| if (image_meter_height == -1): | |
| print("No metadata found for file " + seq_name) | |
| return | |
| # create dataloader | |
| dataloader = create_dataloader_frames_only(in_dir) | |
| try: | |
| inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose) | |
| except: | |
| print("Error in " + seq_name) | |
| with open(os.path.join(out_dir, "ERROR_" + seq_name + ".txt"), 'w') as f: | |
| f.write("ERROR") | |
| return | |
| real_width = image_shapes[0][0][0][1] | |
| real_height = image_shapes[0][0][0][0] | |
| results = do_full_tracking(inference, image_shapes, image_meter_width, image_meter_height, width, height, config=config, gp=None, verbose=verbose) | |
| mot_rows = [] | |
| for frame in results['frames']: | |
| for fish in frame['fish']: | |
| bbox = fish['bbox'] | |
| row = [] | |
| right = bbox[0]*real_width | |
| top = bbox[1]*real_height | |
| w = bbox[2]*real_width - bbox[0]*real_width | |
| h = bbox[3]*real_height - bbox[1]*real_height | |
| row.append(str(frame['frame_num'] + 1)) | |
| row.append(str(fish['fish_id'] + 1)) | |
| row.append(str(int(right))) | |
| row.append(str(int(top))) | |
| row.append(str(int(w))) | |
| row.append(str(int(h))) | |
| row.append("-1") | |
| row.append("-1") | |
| row.append("-1") | |
| row.append("-1") | |
| mot_rows.append(",".join(row)) | |
| mot_text = "\n".join(mot_rows) | |
| with open(os.path.join(out_dir, seq_name + ".txt"), 'w') as f: | |
| f.write(mot_text) | |
| return | |
| def argument_parser(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--frames", required=True, help="Path to frame directory. Required.") | |
| parser.add_argument("--location", required=True, help="Name of location dir. Required.") | |
| parser.add_argument("--metadata", required=True, help="Path to metadata directory. Required.") | |
| parser.add_argument("--output", required=True, help="Path to output directory. Required.") | |
| parser.add_argument("--weights", default='models/v5m_896_300best.pt', help="Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt") | |
| return parser | |
| if __name__ == "__main__": | |
| args = argument_parser().parse_args() | |
| main(args) |