Spaces:
Runtime error
Runtime error
| import os | |
| import torch | |
| from zipfile import ZipFile | |
| from aris import create_manual_marking, BEAM_WIDTH_DIR, create_metadata_dictionary, prep_for_mm | |
| from dataloader import create_dataloader_aris | |
| from inference import do_full_inference, json_dump_round_float | |
| from visualizer import generate_video_batches | |
| WEIGHTS = 'models/v5m_896_300best.pt' | |
| def predict_task(filepath, weights=WEIGHTS, gradio_progress=None): | |
| """ | |
| Main processing task to be run in gradio | |
| - Writes aris frames to dirname(filepath)/frames/{i}.jpg | |
| - Writes json output to dirname(filepath)/{filename}_results.json | |
| - Writes manual marking to dirname(filepath)/{filename}_marking.txt | |
| - Writes video output to dirname(filepath)/{filename}_results.mp4 | |
| - Zips all results to dirname(filepath)/{filename}_results.zip | |
| Args: | |
| filepath (str): path to aris file | |
| TODO: Separate into subtasks in different queues; have a GPU-only queue. | |
| """ | |
| if (gradio_progress): gradio_progress(0, "In task...") | |
| print("Cuda available in task?", torch.cuda.is_available()) | |
| print(filepath) | |
| dirname = os.path.dirname(filepath) | |
| filename = os.path.basename(filepath).replace(".aris","").replace(".ddf","") | |
| results_filepath = os.path.join(dirname, f"{filename}_results.json") | |
| marking_filepath = os.path.join(dirname, f"{filename}_marking.txt") | |
| video_filepath = os.path.join(dirname, f"{filename}_results.mp4") | |
| zip_filepath = os.path.join(dirname, f"{filename}_results.zip") | |
| os.makedirs(dirname, exist_ok=True) | |
| # create dataloader | |
| if (gradio_progress): gradio_progress(0, "Initializing Dataloader...") | |
| dataloader, dataset = create_dataloader_aris(filepath, BEAM_WIDTH_DIR, None) | |
| # extract aris/didson info. didson does not yet have pixel-meter info | |
| if ".ddf" in filepath: | |
| image_meter_width = -1 | |
| image_meter_height = -1 | |
| else: | |
| image_meter_width = dataset.didson.info['xdim'] * dataset.didson.info['pixel_meter_width'] | |
| image_meter_height = dataset.didson.info['ydim'] * dataset.didson.info['pixel_meter_height'] | |
| frame_rate = dataset.didson.info['framerate'] | |
| # run detection + tracking | |
| results = do_full_inference(dataloader, image_meter_width, image_meter_height, gp=gradio_progress, weights=weights) | |
| # re-index results if desired - this should be done before writing the file | |
| results = prep_for_mm(results) | |
| # write output to disk | |
| json_dump_round_float(results, results_filepath) | |
| metadata = None | |
| if dataset.didson.info['version'][3] == 5: # ARIS only | |
| metadata = create_metadata_dictionary(filepath, results_filepath) | |
| create_manual_marking(metadata, out_path=marking_filepath) | |
| # generate a video with tracking results | |
| generate_video_batches(dataset.didson, results_filepath, frame_rate, video_filepath, | |
| image_meter_width=image_meter_width, image_meter_height=image_meter_height, gp=gradio_progress) | |
| # zip up the results | |
| with ZipFile(zip_filepath, 'w') as z: | |
| for file in [results_filepath, marking_filepath, video_filepath, os.path.join(dirname, 'bg_start.jpg')]: | |
| if os.path.exists(file): | |
| z.write(file, arcname=os.path.basename(file)) | |
| # release GPU memory | |
| torch.cuda.empty_cache() | |
| return metadata, results_filepath, zip_filepath, video_filepath, marking_filepath |