get frame from video and generate yolo annotation
import cv2
import os
import csv
import xml.dom.minidom as Document
import xml.dom.minidom as parseString
import pandas as pd
import numpy as np
import dicttoxml
class Size():
def __init__(self, height, width, depth):
self.height = height
self.width = width
self.depth = depth
class Rectangle():
def __init__(self, xmin, ymin, xmax, ymax):
self.xmin = xmin
self.ymin = ymin
self.xmax = xmax
self.ymax = ymax
class Annotation:
def __init__(self, size, rect):
self.folder = ''
self.filename = ''
self.path = ''
self.name = ''
self.size = size
self.rect = rect
def getXML(self, file):
dict = {
"annotation":{
"folder":self.folder,
"filename":self.filename,
"path":self.path,
"source":{
"database":"Unknown"
},
"size":{
"width":self.size.width,
"height":self.size.height,
"depth":3
},
"segmented":0,
"object":{
"name":self.name,
"pose":"Unspecified",
"truncated":0,
"difficult":0,
"bndbox":{
"xmin":self.rect.xmin,
"ymin":self.rect.ymin,
"xmax":self.rect.xmax,
"ymax":self.rect.ymax
}
}
}
}
xml = dicttoxml.dicttoxml(dict, attr_type=False, root=False)
dom = parseString.parseString(xml)
# print(dom.toprettyxml())
f = open(file, 'w', encoding='utf-8')
dom.writexml(f, addindent='\t', newl='\n',encoding='utf-8')
f.close()
def getRect(path, type):
dir = os.path.basename(path)
classes = []
videos = ["hex_screwdriver.mp4"]
for video in videos:
# add class name to classes
class_name = os.path.basename(video)[:-4]
if class_name not in classes:
classes.append(class_name)
# start capture frame from video
cam = cv2.VideoCapture(video)
currentframe = 1
# while(True):
while(currentframe<100):
# reading from frame
ret, frame = cam.read()
if ret:
# if video is still left continue creating images
filename = class_name + "_" + str(currentframe) + '.jpg'
print ('Creating...' + filename)
# produce rectangle from picture
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# gray = cv2.GaussianBlur(gray,(7,7),0)
# _, gray = cv2.threshold(gray,96,255,cv2.THRESH_BINARY_INV)
# delete small points
kernel = np.ones((3, 3), np.uint8)
cv2.erode(gray, kernel)
cv2.dilate(gray, kernel)
x, y, w, h = cv2.boundingRect(gray)
cv2.rectangle(frame, (x,y), (x+w,y+h), (0,255,0), 2)
# writing the extracted images
cv2.imwrite(filename, gray)
# increasing counter
currentframe += 1
size = Size(frame.shape[0], frame.shape[1], frame.shape[2])
if type=="voc":
rect = Rectangle(x, y, x+w, y+h)
annotation = Annotation(size, rect)
annotation.folder = str(dir)
annotation.filename = filename
annotation.path = path + '\\' + filename
annotation.size = size
annotation.name = 'pen'
annotation.rect = rect
annotation.getXML(filename[:-4] + ".xml")
elif type=="yolo":
lst = str(classes.index(class_name))+" "+str((x+w)/size.width)+ \
" "+str((y+h)/size.height)+" "+str(w/size.width)+" "+str(h/size.height)
with open(filename[:-4] + ".txt",'w') as f:
f.write(lst)
else:
break
# Release all space and windows once done
cam.release()
cv2.destroyAllWindows()
with open("classes.txt",'w') as f:
for classname in classes:
f.write(classname+"\n")
getRect(os.path.dirname(os.path.realpath(__file__)), "yolo")
change yolo annotation
import os
import csv
import re
import xml.etree.ElementTree as et
import pandas as pd
import numpy as np
path = "/home/coding/dataset/mmdet-image"
files= os.listdir(path)
for file in files:
# print(file)
# print(re.match('IMG_[0-9]+_[0-9]+.txt', file))
if re.search('IMG_[0-9]+_[0-9]+.txt', file):
full_path = "%s/%s"%(path, file)
# print(full_path)
f = open(full_path,'r')
str = f.read()
f.close()
f = open(full_path,'w')
rows = str.splitlines()
for row in rows:
lst = row.split(" ")
class_number = int(lst[0]) -15
lst[0] = class_number
f.write("%s %s %s %s %s\n"%(lst[0], lst[1], lst[2], lst[3], lst[4]))
# print("%s %s %s %s %s"%(lst[0], lst[1], lst[2], lst[3], lst[4]))
f.close()