Spaces:
Runtime error
Runtime error
File size: 6,286 Bytes
fbb3995 8ca2651 fbb3995 193f172 fbb3995 193f172 fbb3995 193f172 fbb3995 193f172 fbb3995 193f172 fbb3995 193f172 fbb3995 193f172 fbb3995 193f172 fbb3995 |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
import project_path
from lib.yolov5.utils.general import clip_boxes, scale_boxes
import argparse
from datetime import datetime
import torch
import os
from dataloader import create_dataloader_frames_only
from inference import setup_model, do_detection, do_suppression, do_confidence_boost, format_predictions, do_tracking
from visualizer import generate_video_batches
import json
from tqdm import tqdm
import numpy as np
def main(args, config={}, verbose=True):
"""
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.
"""
print("In task...")
print("Cuda available in task?", torch.cuda.is_available())
# setup config
if "conf_threshold" not in config: config['conf_threshold'] = 0.001
if "nms_iou" not in config: config['nms_iou'] = 0.6
if "min_length" not in config: config['min_length'] = 0.3
if "max_age" not in config: config['max_age'] = 20
if "iou_threshold" not in config: config['iou_threshold'] = 0.01
if "min_hits" not in config: config['min_hits'] = 11
print(config)
dirname = args.frames
locations = [
"elwha-val"
]
for loc in locations:
in_loc_dir = os.path.join(dirname, loc)
out_dir = os.path.join(args.output, loc, "tracker", "data")
print(in_loc_dir)
print(out_dir)
# run detection + tracking
seq_list = os.listdir(in_loc_dir)
iterate_sequences(in_loc_dir, out_dir, config, args.weights, seq_list, verbose)
def iterate_sequences(in_dir, out_dir, config, weights, seq_list, verbose):
model, device = setup_model(weights)
idx = 1
with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
for seq in seq_list:
ann_list = []
pbar.update(1)
if (seq.startswith(".")): continue
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_dir, seq)
frame_list = detect(in_seq_dir, out_dir, config, seq, model, device, verbose)
for frame in frame_list:
if frame is not None:
for ann in frame:
ann_list.append({
'image_id': ann[5],
'category_id': 0,
'bbox': [ann[0], ann[1], ann[2] - ann[0], ann[3] - ann[1]],
'score': ann[4]
})
result = json.dumps(ann_list)
out_seq_dir = os.path.join(out_dir, seq)
os.makedirs(out_seq_dir, exist_ok=True)
with open(os.path.join(out_seq_dir, 'pred.json'), 'w') as f:
f.write(result)
def iterate_files(in_dir, out_dir, config, weights, verbose):
model, device = setup_model(weights)
ann_list = []
frame_list = detect(in_dir, out_dir, config, "specified folder", model, device, verbose)
with tqdm(total=len(frame_list), desc="...", ncols=0) as pbar:
for frame in frame_list:
if frame is not None:
for ann in frame:
ann_list.append({
'image_id': ann[5],
'category_id': 0,
'bbox': [ann[0], ann[1], ann[2] - ann[0], ann[3] - ann[1]],
'score': ann[4]
})
pbar.update(1)
result = json.dumps(ann_list)
with open(os.path.join(out_dir, 'pred.json'), 'w') as f:
f.write(result)
def detect(in_dir, out_dir, config, seq_name, model, device, verbose):
#progress_log = lambda p, m: 0
# 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
outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
file_names = dataloader.files
frame_list = []
for batch_i, batch in enumerate(outputs):
batch_shapes = image_shapes[batch_i]
# Format results
for si, pred in enumerate(batch):
(image_shape, original_shape) = batch_shapes[si]
# Clip boxes to image bounds and resize to input shape
clip_boxes(pred, (height, width))
boxes = pred[:, :4].clone() # xyxy
confs = pred[:, 4].clone().tolist()
scale_boxes(image_shape, boxes, original_shape[0], original_shape[1]) # to original shape
frame = [ [*bb, conf, file_name] for bb, conf, file_name in zip(boxes.tolist(), confs, file_names[batch_i*32+si]) ]
file_name = file_names[batch_i*32 + si]
for ann in frame:
ann.append(file_name)
frame_list.append(frame)
return frame_list
def argument_parser():
parser = argparse.ArgumentParser()
parser.add_argument("--frames", required=True, help="Path to frame 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) |