diff --git a/src/main.zig b/src/main.zig index 2232e3b..1d2cae5 100644 --- a/src/main.zig +++ b/src/main.zig @@ -22,9 +22,15 @@ const max_edge_layers = 32; var exit_loop: bool = false; +const CliArgs = struct { + hef_file_path: []const u8, + gst_source_path: []const u8, +}; pub fn main() !void { + var args: CliArgs = .{ .hef_file_path = undefined, .gst_source_path = undefined }; + std.posix.sigaction(std.posix.SIG.INT, &std.posix.Sigaction { .handler = .{ .handler = struct { @@ -45,13 +51,50 @@ pub fn main() !void { // This allows me to utilize the same command line args and gstreamer gst.init(@ptrCast(&std.os.argv.len), @ptrCast(&std.os.argv.ptr)); - var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator); + var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator); defer arena.deinit(); const alloc = arena.allocator(); + + var arg_source = try std.process.argsWithAllocator(alloc); + _ = arg_source.skip(); + + var hef_set = false; + var gst_set = false; + + + while (arg_source.next()) |cur_arg| { + if (std.mem.eql(u8, cur_arg, "--hef")) { + + if (arg_source.next()) |x| { + args.hef_file_path = x; + hef_set = true; + } else { + std.process.fatal("--hef flag specified, but no value provided!\n", .{}); + } + } + if (std.mem.eql(u8, cur_arg, "--gst-source")) { + if (arg_source.next()) |x| { + args.gst_source_path = x; + gst_set = true; + } else { + std.process.fatal("--gst-source flag specified, but no value provided!\n", .{}); + } + } + } + + if (!gst_set) { + std.process.fatal("Invalid usage. --gst-source value is required\n", .{}); + } + if (!hef_set) { + std.process.fatal("Invalid usage. --hef value is required\n", .{}); + } + + + std.debug.print("Using hardcoded hef file path :P\n", .{}); std.fs.cwd().access(hef_file, .{ }) catch |e| { - std.debug.panic("Could not open hef file! '{any}'", .{ e }); + std.process.fatal("Could not open hef file! '{any}'", .{ e }); }; var status: hlo.hailo_status = undefined; @@ -119,129 +162,133 @@ pub fn main() !void { std.debug.print("HailoRT init completed\nGstreamer Init starting", .{}); - std.debug.print("Gstreamer Initialized\n", .{}); + _ = gst_block: { + std.debug.print("Gstreamer Initialized\n", .{}); - const source: *gst.Element = gst.ElementFactory.make("v4l2src", "source") orelse unreachable; - std.debug.print("test source created\n", .{}); - const scale: *gst.Element = gst.ElementFactory.make("videoscale", "scale") orelse unreachable; - std.debug.print("video scale created\n", .{}); - const format: *gst.Element = gst.ElementFactory.make("videoconvert", "format") orelse unreachable; - std.debug.print("video convert created\n", .{}); - const sink_el: *gst.Element = gst.ElementFactory.make("appsink", "sink") orelse unreachable; - std.debug.print("appsink created\n", .{}); - const sink: *gstapp.AppSink = gobject.ext.cast(gstapp.AppSink, sink_el) orelse unreachable; - std.debug.print("appsink cast to an AppSink\n", .{}); + const source: *gst.Element = gst.ElementFactory.make("v4l2src", "source") orelse unreachable; + std.debug.print("test source created\n", .{}); + const scale: *gst.Element = gst.ElementFactory.make("videoscale", "scale") orelse unreachable; + std.debug.print("video scale created\n", .{}); + const format: *gst.Element = gst.ElementFactory.make("videoconvert", "format") orelse unreachable; + std.debug.print("video convert created\n", .{}); + const sink_el: *gst.Element = gst.ElementFactory.make("appsink", "sink") orelse unreachable; + std.debug.print("appsink created\n", .{}); + const sink: *gstapp.AppSink = gobject.ext.cast(gstapp.AppSink, sink_el) orelse unreachable; + std.debug.print("appsink cast to an AppSink\n", .{}); - const sink_caps = gst.Caps.fromString("video/x-raw,format=RGB,width=640,height=640") orelse unreachable; - std.debug.print("sink caps created\n", .{}); - sink.setCaps(sink_caps); - std.debug.print("sink caps applied\n", .{}); + const sink_caps = gst.Caps.fromString("video/x-raw,format=RGB,width=640,height=640") orelse unreachable; + std.debug.print("sink caps created\n", .{}); + sink.setCaps(sink_caps); + std.debug.print("sink caps applied\n", .{}); - const pipeline: *gst.Pipeline = gst.Pipeline.new("test-pipeline"); - std.debug.print("pipeline created\n", .{}); + const pipeline: *gst.Pipeline = gst.Pipeline.new("test-pipeline"); + std.debug.print("pipeline created\n", .{}); - const bin: *gst.Bin = &pipeline.f_bin; - std.debug.print("Bin retrieved from Pipeline\n", .{}); + const bin: *gst.Bin = &pipeline.f_bin; + std.debug.print("Bin retrieved from Pipeline\n", .{}); - _ = bin.add(source); - std.debug.print("source added to bin\n", .{}); - _ = bin.add(scale); - std.debug.print("scale added to bin\n", .{}); - _ = bin.add(format); - std.debug.print("format added to bin\n", .{}); - _ = bin.add(sink_el); - // _ = bin.addMany(source, scale, format, sink_el); - std.debug.print("Elements added to bin\n", .{}); + _ = bin.add(source); + std.debug.print("source added to bin\n", .{}); + _ = bin.add(scale); + std.debug.print("scale added to bin\n", .{}); + _ = bin.add(format); + std.debug.print("format added to bin\n", .{}); + _ = bin.add(sink_el); + // _ = bin.addMany(source, scale, format, sink_el); + std.debug.print("Elements added to bin\n", .{}); - // the failure return code is -1 I believe - if (gst.Element.linkMany(source, scale, format, sink_el) < 0) { - pipeline.unref(); - std.debug.panic("Elements could not be linked\n", .{}); - } - std.debug.print("Elements linked\n", .{}); - - // g_int is just i32. You can - // source.set("pattern", @as(i16, 0)); - // gobject.Object.set(source.as(gobject.Object), "pattern", @as(i16, 0)); - - const ret = pipeline.as(gst.Element).setState(gst.State.playing); - if (ret == gst.StateChangeReturn.failure) { - pipeline.unref(); - std.debug.panic("Could not start pipeline", .{}); - } - - const bus: *gst.Bus = pipeline.getBus(); - - - std.debug.print("Getting vstream info\n", .{}); - var stream_info: hlo.hailo_vstream_info_t = undefined; - status = hlo.hailo_get_input_vstream_info(input_vstreams[0], &stream_info); - assert(status == hlo.HAILO_SUCCESS); - - var output_info: hlo.hailo_vstream_info_t = undefined; - status = hlo.hailo_get_output_vstream_info(output_vstreams[0], &output_info); - assert(status == hlo.HAILO_SUCCESS); - std.debug.print("\nOutput info: {any}\n\n", .{ output_info }); - - - var input_frame_size: usize = 0; - - status = hlo.hailo_get_input_vstream_frame_size(input_vstreams[0], &input_frame_size); - assert(status == hlo.HAILO_SUCCESS); - - const frame_count = input_frame_size / @sizeOf(u8); - const input_data: [:0]u8 = try alloc.allocSentinel(u8, frame_count, 0); - - - var output_frame_size: usize = 0; - status = hlo.hailo_get_output_vstream_frame_size(output_vstreams[0], &output_frame_size); - assert(status == hlo.HAILO_SUCCESS); - std.debug.print("output frame size is: {d}\n", .{ output_frame_size }); - - const output_frame_count = output_frame_size / @sizeOf(f32); - const output_data: [:0]f32 = try alloc.allocSentinel(f32, output_frame_count, 0); - - - var run_count: u8 = 0; - - while (!exit_loop) { - const sample = sink.pullSample() orelse unreachable; - const buffer: *gst.Buffer = sample.getBuffer() orelse unreachable; - - _ = buffer.extract(0, @ptrCast(input_data), input_frame_size); - - status = hlo.hailo_vstream_write_raw_buffer(input_vstreams[0], input_data.ptr, input_frame_size); - assert(status == hlo.HAILO_SUCCESS); - - status = hlo.hailo_flush_input_vstream(input_vstreams[0]); - assert(status == hlo.HAILO_SUCCESS); - - status = hlo.hailo_vstream_read_raw_buffer(output_vstreams[0], output_data.ptr, output_frame_size); - if (status == hlo.HAILO_SUCCESS) { - for (0..50) |x| { - std.debug.print(" {d} ", .{output_data[x]}); - } - std.debug.print("\n", .{}); - - run_count += 1; - if (run_count > 20) { - break; - } - } else { - std.debug.print("Read failed!\n", .{}); - continue; + // the failure return code is -1 I believe + if (gst.Element.linkMany(source, scale, format, sink_el) < 0) { + pipeline.unref(); + std.debug.print("Elements could not be linked! Aborting!\n", .{}); + break :gst_block; } - } + std.debug.print("Elements linked\n", .{}); - std.debug.print("Staring shutdown\n", .{ }); + // g_int is just i32. You can + // source.set("pattern", @as(i16, 0)); + // gobject.Object.set(source.as(gobject.Object), "pattern", @as(i16, 0)); + + const ret = pipeline.as(gst.Element).setState(gst.State.playing); + if (ret == gst.StateChangeReturn.failure) { + pipeline.unref(); + std.debug.print("Could not start pipeline! Aboring\n", .{}); + break :gst_block; + } + + const bus: *gst.Bus = pipeline.getBus(); - bus.unref(); - _ = pipeline.as(gst.Element).setState(gst.State.null); - pipeline.unref(); - // const bus: *gst.Bus = pipeline.getBus(); + std.debug.print("Getting vstream info\n", .{}); + var stream_info: hlo.hailo_vstream_info_t = undefined; + status = hlo.hailo_get_input_vstream_info(input_vstreams[0], &stream_info); + assert(status == hlo.HAILO_SUCCESS); + + var output_info: hlo.hailo_vstream_info_t = undefined; + status = hlo.hailo_get_output_vstream_info(output_vstreams[0], &output_info); + assert(status == hlo.HAILO_SUCCESS); + std.debug.print("\nOutput info: {any}\n\n", .{ output_info }); + + + var input_frame_size: usize = 0; + + status = hlo.hailo_get_input_vstream_frame_size(input_vstreams[0], &input_frame_size); + assert(status == hlo.HAILO_SUCCESS); + + const frame_count = input_frame_size / @sizeOf(u8); + const input_data: [:0]u8 = try alloc.allocSentinel(u8, frame_count, 0); + + + var output_frame_size: usize = 0; + status = hlo.hailo_get_output_vstream_frame_size(output_vstreams[0], &output_frame_size); + assert(status == hlo.HAILO_SUCCESS); + std.debug.print("output frame size is: {d}\n", .{ output_frame_size }); + + const output_frame_count = output_frame_size / @sizeOf(f32); + const output_data: [:0]f32 = try alloc.allocSentinel(f32, output_frame_count, 0); + + defer pipeline.unref(); + defer _ = pipeline.as(gst.Element).setState(gst.State.null); + defer bus.unref(); + + var run_count: u8 = 0; + + while (!exit_loop) { + const sample = sink.pullSample() orelse unreachable; + const buffer: *gst.Buffer = sample.getBuffer() orelse unreachable; + + _ = buffer.extract(0, @ptrCast(input_data), input_frame_size); + + status = hlo.hailo_vstream_write_raw_buffer(input_vstreams[0], input_data.ptr, input_frame_size); + assert(status == hlo.HAILO_SUCCESS); + + status = hlo.hailo_flush_input_vstream(input_vstreams[0]); + assert(status == hlo.HAILO_SUCCESS); + + status = hlo.hailo_vstream_read_raw_buffer(output_vstreams[0], output_data.ptr, output_frame_size); + if (status == hlo.HAILO_SUCCESS) { + for (0..(output_data.len / 5)) |x| { + if (output_data[x] == 1.0) { + + std.debug.print("Person found! Confidence: {d}, {d}:{d}, {d}, {d}\n", .{ output_data[x + 1], output_data[x + 2], output_data[x + 3], output_data[x + 4], output_data[x + 5] }); + } + } + + run_count += 1; + if (run_count > 20) { + break; + } + } else { + std.debug.print("Read failed!\n", .{}); + continue; + } + } + + std.debug.print("Staring shutdown\n", .{ }); + + }; _ = hlo.hailo_release_output_vstreams(&output_vstreams, output_vstream_size);