11import unittest
22import sys
33import os
4+ from lxml import etree
5+ import pandas as pd
6+ import cv2
47from kabr_tools import cvat2ultralytics
58from tests .utils import (
69 del_dir ,
710 del_file ,
8- get_detection
11+ get_detection ,
12+ dir_exists ,
13+ file_exists
914)
1015
1116
@@ -32,7 +37,7 @@ def setUp(self):
3237 self .video = TestCvat2Ultralytics .dir
3338 self .annotation = TestCvat2Ultralytics .dir
3439 self .dataset = "tests/ultralytics"
35- self .skip = "5 "
40+ self .skip = "1 "
3641 self .label2index = "ethogram/label2index.json"
3742
3843 def tearDown (self ):
@@ -44,9 +49,75 @@ def test_run(self):
4449 sys .argv = [self .tool ,
4550 "--video" , self .video ,
4651 "--annotation" , self .annotation ,
47- "--dataset" , "tests/ultralytics" ]
52+ "--dataset" , "tests/ultralytics" ,
53+ "--skip" , self .skip ]
4854 run ()
4955
56+ # check for output dirs
57+ self .assertTrue (dir_exists (self .dataset ))
58+ self .assertTrue (dir_exists (f"{ self .dataset } ))
59+ self .assertTrue (dir_exists (f"{ self .dataset } ))
60+ self .assertTrue (dir_exists (f"{ self .dataset } ))
61+ self .assertTrue (dir_exists (f"{ self .dataset } ))
62+ self .assertTrue (dir_exists (f"{ self .dataset } ))
63+ self .assertTrue (dir_exists (f"{ self .dataset } ))
64+
65+ # check output
66+ annotations = etree .parse (TestCvat2Ultralytics .annotation ).getroot ()
67+ tracks = [list (track .findall ("box" )) for track in annotations .findall ("track" )]
68+ self .assertEqual (len (tracks [0 ]), 21 )
69+ self .assertEqual (len (tracks [0 ]), len (tracks [1 ]))
70+ original_size = annotations .find ("meta" ).find ("task" ).find ("original_size" )
71+ height = int (original_size .find ("height" ).text )
72+ width = int (original_size .find ("width" ).text )
73+ for i in range (len (tracks [0 ])):
74+ # check existence
75+ if i < 16 :
76+ data_im = f"{ self .dataset } { i }
77+ self .assertTrue (file_exists (data_im ))
78+ data_label = f"{ self .dataset } { i }
79+ self .assertTrue (file_exists (data_label ))
80+ elif i < 18 :
81+ data_im = f"{ self .dataset } { i }
82+ self .assertTrue (file_exists (data_im ))
83+ data_label = f"{ self .dataset } { i }
84+ self .assertTrue (file_exists (data_label ))
85+ else :
86+ data_im = f"{ self .dataset } { i }
87+ self .assertTrue (file_exists (data_im ))
88+ data_label = f"{ self .dataset } { i }
89+ self .assertTrue (file_exists (data_label ))
90+
91+ # check image
92+ data_im = cv2 .imread (data_im )
93+ self .assertEqual (data_im .shape , (height , width , 3 ))
94+
95+ # check label
96+ data_label = pd .read_csv (data_label , sep = " " , header = None )
97+ annotation_label = []
98+ for track in tracks :
99+ box = track [i ]
100+ x_start = float (box .attrib ["xtl" ])
101+ y_start = float (box .attrib ["ytl" ])
102+ x_end = float (box .attrib ["xbr" ])
103+ y_end = float (box .attrib ["ybr" ])
104+ x_center = (x_start + (x_end - x_start ) / 2 ) / width
105+ y_center = (y_start + (y_end - y_start ) / 2 ) / height
106+ w = (x_end - x_start ) / width
107+ h = (y_end - y_start ) / height
108+ annotation_label .append (
109+ [0 , x_center , y_center , w , h ]
110+ )
111+ self .assertEqual (len (data_label .index ), len (annotation_label ))
112+
113+ for i , label in enumerate (annotation_label ):
114+ self .assertEqual (label [0 ], annotation_label [i ][0 ])
115+ self .assertAlmostEqual (label [1 ], annotation_label [i ][1 ], places = 4 )
116+ self .assertAlmostEqual (label [2 ], annotation_label [i ][2 ], places = 4 )
117+ self .assertAlmostEqual (label [3 ], annotation_label [i ][3 ], places = 4 )
118+ self .assertAlmostEqual (label [4 ], annotation_label [i ][4 ], places = 4 )
119+
120+
50121 def test_parse_arg_min (self ):
51122 # parse arguments
52123 sys .argv = [self .tool ,
@@ -81,7 +152,7 @@ def test_parse_arg_full(self):
81152 self .assertEqual (args .video , self .video )
82153 self .assertEqual (args .annotation , self .annotation )
83154 self .assertEqual (args .dataset , self .dataset )
84- self .assertEqual (args .skip , 5 )
155+ self .assertEqual (args .skip , int ( self . skip ) )
85156 self .assertEqual (args .label2index , self .label2index )
86157
87158 # run cvat2ultralytics
0 commit comments