Convert a list of basic CSV annotations into JSON files and split the corresponding PNGs randomly into train, test, val folders

to_coco

#!/bin/ruby

require 'json'
require 'csv'
require 'fileutils'

# Run this script in a directory with all PNGs and labels in COCO format in the annotations folder

folders = [ 'train', 'test', 'val' ]
proba = [ 0.6, 0.2, 0.2 ]

# create folders if they don't
folders.each do |folder|
	Dir.mkdir(folder) if not Dir.exists?(folder)
end

# get list of all png files
files = Dir['*.png', 'test/*.png', 'train/*.png', 'val/*.png'].shuffle
file_number = files.length
proba.map!{|x| (x*file_number).round} # how many images in each category

# sort all files
files.each_with_index do |file, index|
	# if the png isn't already sorted
	if not folders.include? File.basename(File.dirname(file))
		# split into 3 folders
		if index < proba[0]
			FileUtils.mv(file, folders[0])
		elsif index < proba[0] + proba[1]
			FileUtils.mv(file, folders[1])
		else
			FileUtils.mv(file, folders[2])
		end
	end
end

# generate jsons
height = 480
width = 640
image_id = 0

folders.each do |folder|
	json = JSON.parse("{}")

	json["images"] = []
	json["annotations"] = []
	json["categories"] = 
	[
		{'id': 0, 'name': 'rock'},
		{'id': 1, 'name': 'slope'},
		{'id': 2, 'name': 'rover'}
	]

	Dir.glob("#{folder}/frame*.png") do |file|
		file_name = File.basename(file)
		image_id = image_id + 1
		json["images"].append({'file_name': file_name, 'height': height, 'width': width, 'id': image_id})

		CSV.foreach("annotations/#{File.basename(file, ".png")}.txt", col_sep:' ') do |row|
			id, category_id, x, y, w, h, _ = row
			# avoid out of bounds coordinates (< 0, > width, height)
			x = x.to_i; y = y.to_i; w = w.to_i; h = h.to_i
			min_x = [0,x].max
			min_y = [0,y].max
			max_x = [width, x + w].min
			max_y = [height, y + h].min
			segmentation = [ max_x, min_y, max_x, max_y, min_x, max_y, min_x, min_y ]
			json["annotations"].append({'image_id': image_id.to_i, 'category_id': category_id.to_i, 'id': id.to_i, 'bbox': [ x,y,w,h ], 'area': w*h, 'iscrowd': false, "isbbox": true, 'segmentation': segmentation})
		end
	end

	File.open("annotations/#{folder}.json","w") do |f|
		f.write(JSON.pretty_generate(json))
	end
end