video.py 6.59 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""

Script to detect yield signs from video.  Videos are loaded into ram,
so you need to split the file or else it starts swapping ;)

"""


import sys
import argparse
14
import configparser
15 16 17

import numpy as np
import skvideo.io as skvio
18
import skimage.draw as skd
19
from tqdm import tqdm
20 21


22
from src import video
23 24
from src import logger
from src import tmeasure
25 26 27 28

# ---


29 30
ROI_LIFESPAN = 30

31
log = logger(name=__name__[2:-2])
32
cfg = configparser.ConfigParser()
33 34 35 36


# ---

37

38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55
def load(fname: str) -> np.ndarray:
    vid = skvio.vread(fname)
    n, h, w, depth = vid.shape
    if depth != 3:
        print('You need to provide a colored video!')
        sys.exit(2)

    log.info('loaded video with %d frames and resulution %d, %d',
             n, h, w)

    return vid


def save(fname: str, vid: np.ndarray):
    log.info('saving to %s', fname)
    skvio.vwrite(fname, vid)


Felix Hamann's avatar
Felix Hamann committed
56
def _draw_vertices(frame, roi):
57
    rr, cc = skd.polygon_perimeter(
Felix Hamann's avatar
Felix Hamann committed
58 59
        roi.vy + (roi.vy[0], ),
        roi.vx + (roi.vx[0], ),
60 61
        shape=frame.shape)

Felix Hamann's avatar
Felix Hamann committed
62
    frame[rr, cc] = [int(255 * roi.health / ROI_LIFESPAN), 0, 0]
63 64


65
def _draw_roi(arr, i, roi, val):
66 67
    y0, y1 = roi.r0 - 1, roi.r1
    x0, x1 = roi.c0 - 1, roi.c1
68

69 70
    # clockwise from topleft
    rr, cc = skd.polygon_perimeter(
71
        (y0, y0, y1, y1, y0),
72 73
        (x0, x1, x1, x0, x0),
        shape=arr[i].shape)
74

75
    arr[i, rr, cc] = val
76

77
    # color indicator for life decay
78
    rr, cc = skd.circle(y0, x0, 5, shape=arr[i].shape)
79
    arr[i, rr, cc] = val * (roi.health / ROI_LIFESPAN)
Felix Hamann's avatar
Felix Hamann committed
80 81
    rr, cc = skd.circle_perimeter(y0, x0, 5)
    arr[i, rr, cc] = val
82 83


84 85 86 87 88 89 90 91 92
def _get_vertices(keys, pois):
    p0, p1, p2 = keys
    vertices = (p0, p1), (p0, p2), (p1, p2)
    vy, vx = zip(*[pois[a][b] for a, b in vertices])
    return vy, vx


def _find_rois(buf, frame, barycenters, pois):
    _, w, h, _ = buf.original.shape
93 94
    rois = []

95 96
    for keys, (y, x) in barycenters:
        vy, vx = _get_vertices(keys, pois)
Felix Hamann's avatar
Felix Hamann committed
97
        rois.append(video.ROI(w, h, y, x, vy, vx, ROI_LIFESPAN))
98 99 100 101

    return rois


102
def _scan_full(buf, frame, i, pipe):
Felix Hamann's avatar
Felix Hamann committed
103 104
    segments = pipe.binarize(frame.astype(np.int64))

105
    buf.binary[i] = segments
106 107
    if pipe.binary and not pipe.edges:
        return []
108

Felix Hamann's avatar
Felix Hamann committed
109 110
    edges = pipe.edge(segments)

111
    buf.edges[i] = edges
Felix Hamann's avatar
Felix Hamann committed
112

113 114
    if pipe.edges:
        return []
115

Felix Hamann's avatar
Felix Hamann committed
116 117 118 119
    barycenters, pois = pipe.detect(edges, segments)
    rois = _find_rois(buf, frame, barycenters, pois)

    return rois
120

121

122 123
def _scan_roi(buf, frame, i, pipe, roi):
    r0, r1, c0, c1 = roi.r0, roi.r1, roi.c0, roi.c1
124

125 126
    buf.binary[i, r0:r1, c0:c1] = pipe.binarize(
        frame[r0:r1, c0:c1].astype(np.int64))
127

128
    _draw_roi(buf.binary, i, roi, 255)
129

130 131
    buf.edges[i, r0:r1, c0:c1] = pipe.edge(
        buf.binary[i, r0:r1, c0:c1])
132

133
    _draw_roi(buf.edges, i, roi, 255)
134

135 136 137
    return pipe.detect(
        buf.edges[i, r0:r1, c0:c1],
        buf.binary[i, r0:r1, c0:c1])
138

139

Felix Hamann's avatar
Felix Hamann committed
140 141 142 143 144 145 146 147 148 149 150
def _merge_rois(w, h, roi, barycenter, pois, new_rois):
    keys, (y, x) = barycenter
    vy, vx = _get_vertices(keys, pois)
    vy = tuple(map(lambda y: roi.r0 + y, vy))
    vx = tuple(map(lambda x: roi.c0 + x, vx))

    new_roi = video.ROI(w, h, y, x, vy, vx, ROI_LIFESPAN)
    if not any([new_roi.intersects(r) for r in new_rois]):
        new_rois.append(new_roi)


151 152 153 154 155
def _process(buf, frame, i, pipe, rois):
    _, w, h, _ = buf.original.shape
    new_rois = []

    # full scan
Felix Hamann's avatar
Felix Hamann committed
156
    rescan = i % 15 == 0
Felix Hamann's avatar
Felix Hamann committed
157 158
    no_rois = len(rois) == 0
    if rescan or no_rois:
159
        new_rois += _scan_full(buf, frame, i, pipe)
160 161 162 163

    # roi scan
    for roi in rois:
        barycenters, pois = _scan_roi(buf, frame, i, pipe, roi)
164
        if len(barycenters) > 0:
Felix Hamann's avatar
Felix Hamann committed
165
            _merge_rois(w, h, roi, list(barycenters)[0], pois, new_rois)
166 167 168 169
        else:
            if not roi.dead:
                roi.punish()
                new_rois.append(roi)
170

Felix Hamann's avatar
Felix Hamann committed
171 172 173
    # nice drawings...
    frame //= 3
    for roi in new_rois:
Felix Hamann's avatar
Felix Hamann committed
174
        _draw_vertices(frame, roi)
Felix Hamann's avatar
Felix Hamann committed
175 176
        _draw_roi(buf.original, i, roi, np.array([255, 255, 255]))

177 178 179 180 181 182
    return new_rois


def process(vid: np.ndarray, only, config=None) -> video.Buffer:
    only_binary, only_edges = only
    only_binary = only_edges or only_binary
183 184 185

    log.info('initializing modules')

186
    pipe = video.Pipeline(config, only_binary, only_edges)
187
    buf = video.Buffer(vid)
188

189
    log.info('start processing')
190 191 192

    n, w, h, _ = vid.shape
    rois = []
193 194 195

    print('')
    done = tmeasure(log.info, 'took %sms')
196
    stats = []
197

198
    for i, frame in tqdm(buf, total=buf.framecount, unit='frames'):
199
        stats.append(len(rois) if len(rois) > 0 else '.')
200
        rois = _process(buf, frame, i, pipe, rois)
201

202
    print(''.join(map(str, stats)) + '\n')
203 204
    done()

205
    return buf
206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222

#
# --- initialization
#


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        'f_in', type=str,
        help='input file')

    parser.add_argument(
        'f_out', type=str,
        help='output file')

223 224 225 226 227 228 229 230 231 232 233 234 235 236
    parser.add_argument(
        '--config', type=str, nargs=1,
        help='configuration file')

    parser.add_argument(
        '--binary',
        action='store_true',
        help='only apply segmentation and morphology')

    parser.add_argument(
        '--edges',
        action='store_true',
        help='only apply --binary and edge detection')

237
    parser.add_argument(
238
        '--save-all',
239 240 241
        action='store_true',
        help='save not only the result but all intermediate steps')

242 243 244 245 246 247 248
    args = parser.parse_args()

    if args.binary and args.edges:
        print('either provide --edges or --binary')
        parser.print_help()
        sys.exit(2)

249 250 251 252 253
    if args.save_all and (args.binary or args.edges):
        print('you can not save_all when using --edges or --binary')
        parser.print_help()
        sys.exit(2)

254
    return args
255 256


257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280
def _save(buf, args):
    if args.save_all:
        out = args.f_out.rsplit('.', maxsplit=1)

        out.insert(-1, 'binary')
        save('.'.join(out), buf.binary)

        out[-2] = 'edges'
        save('.'.join(out), buf.edges)

        save(args.f_out, buf.original)
        return

    if args.edges:
        save(args.f_out, buf.edges)
        return

    if args.binary:
        save(args.f_out, buf.binary)
        return

    save(args.f_out, buf.original)


281 282
def main(args):
    log.info('starting the application')
283 284 285 286 287 288 289

    if len(args.config) > 0:
        cfg.read(args.config[0])
        if 'options' not in cfg:
            print('You need to provide an [option] section')
            sys.exit(2)

290 291 292
    buf = process(load(args.f_in),
                  (args.binary, args.edges),
                  config=cfg if args.config else None)
293

294
    _save(buf, args)
295 296 297 298 299


if __name__ == '__main__':
    args = parse_args()
    main(args)