Spaces:
Runtime error
Runtime error
File size: 3,438 Bytes
7a4b92f 5ab0373 7a4b92f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
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 |