extract_red_contours.py

#!/usr/local/bin/python3
import glob
import numpy as np
import cv2
import sys
import os
import shutil

label='Ossicle'

if __name__ == '__main__':

	img_file_path_list = glob.glob('original/*.png')

	if not os.path.exists('./images'):
		os.mkdir('images')

	if not os.path.exists('./labels'):
		os.mkdir('labels')

	for i, img_file_path in enumerate(img_file_path_list):

		img = cv2.imread(img_file_path, cv2.IMREAD_COLOR)
		hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

		# Detect red rectangle 
		lower_red = np.array([0, 255, 255])
		upper_red = np.array([1, 255, 255])
		img_mask = cv2.inRange(hsv, lower_red, upper_red)

		# Find contours
		image, contours, hierarchy = cv2.findContours(img_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

		# Copy an original image to images folder
		file_num = '%05d' % i
		shutil.copyfile(img_file_path, 'images/'+str(file_num)+'.png')

		# Create a label file
		label_file_path = 'labels/'+str(file_num)+'.txt'
		label_file = open(label_file_path, 'w')

		print(label_file_path)

		for X in contours:
			left = X[:,0][0][0]
			top = X[:,0][0][1]
			right = X[:,0][2][0]
			bottom = X[:,0][2][1]

			elements = (label, 0.00, 0, -10, left, top, right, bottom, -1, -1, -1, -1000, -1000, -1000, -10)
			line = ' '.join(str(x) for x in elements)
			label_file.write(line+'\n')

			print(line)