66import os
77
88import cv2
9- import datumaro as dm
109import numpy as np
10+ from datumaro .components .annotation import Label , Bbox , Mask
11+ from datumaro .components .dataset import Dataset
12+ from datumaro .components .dataset_base import DatasetItem
13+ from datumaro .components .media import ImageFromFile , ImageFromNumpy
14+
1115
1216from otx .api .entities .model_template import TaskType
1317
@@ -107,7 +111,7 @@ def generate_datumaro_dataset_item(
107111 image_shape : np .array = np .array ((5 , 5 , 3 )),
108112 mask_shape : np .array = np .array ((5 , 5 )),
109113 temp_dir : Optional [str ] = None ,
110- ) -> dm . DatasetItem :
114+ ) -> DatasetItem :
111115 """Generate Datumaro DatasetItem.
112116
113117 Args:
@@ -119,20 +123,22 @@ def generate_datumaro_dataset_item(
119123 temp_dir (str): directory to save image data
120124
121125 Returns:
122- dm. DatasetItem: Datumaro DatasetItem
126+ DatasetItem: Datumaro DatasetItem
123127 """
124128 ann_task_dict = {
125- "classification" : dm . Label (label = 0 ),
126- "detection" : dm . Bbox (1 , 2 , 3 , 4 , label = 0 ),
127- "segmentation" : dm . Mask (np .zeros (mask_shape )),
129+ "classification" : Label (label = 0 ),
130+ "detection" : Bbox (1 , 2 , 3 , 4 , label = 0 ),
131+ "segmentation" : Mask (np .zeros (mask_shape )),
128132 }
129133
130134 if temp_dir :
131135 path = os .path .join (temp_dir , "image.png" )
132136 cv2 .imwrite (path , np .ones (image_shape ))
133- return dm . DatasetItem (id = item_id , subset = subset , image = path , annotations = [ann_task_dict [task ]])
137+ return DatasetItem (id = item_id , subset = subset , media = ImageFromFile ( path ) , annotations = [ann_task_dict [task ]])
134138
135- return dm .DatasetItem (id = item_id , subset = subset , image = np .ones (image_shape ), annotations = [ann_task_dict [task ]])
139+ return DatasetItem (
140+ id = item_id , subset = subset , media = ImageFromNumpy (np .ones (image_shape )), annotations = [ann_task_dict [task ]]
141+ )
136142
137143
138144def generate_datumaro_dataset (
@@ -141,7 +147,7 @@ def generate_datumaro_dataset(
141147 num_data : int = 1 ,
142148 image_shape : np .array = np .array ((5 , 5 , 3 )),
143149 mask_shape : np .array = np .array ((5 , 5 )),
144- ) -> dm . Dataset :
150+ ) -> Dataset :
145151 """Generate Datumaro Dataset.
146152
147153 Args:
@@ -154,7 +160,7 @@ def generate_datumaro_dataset(
154160 Returns:
155161 dm.Dataset: Datumaro Dataset
156162 """
157- dataset_items : dm . DatasetItem = []
163+ dataset_items : DatasetItem = []
158164 for subset in subsets :
159165 for idx in range (num_data ):
160166 dataset_items .append (
@@ -166,4 +172,4 @@ def generate_datumaro_dataset(
166172 mask_shape = mask_shape ,
167173 )
168174 )
169- return dm . Dataset .from_iterable (dataset_items , categories = ["cat" , "dog" ])
175+ return Dataset .from_iterable (dataset_items , categories = ["cat" , "dog" ])
0 commit comments