Coverage for collectra/converters/yolo.py: 81.73%
104 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1import shutil
2from pathlib import Path
4from pydantic import BaseModel, Field
6from ..commons.files import CollectraFile
7from .base import Converter
10class YOLOConfig(BaseModel):
11 """
12 Configuration for YOLO dataset conversion.
13 This class defines the structure of the YOLO configuration file, which includes:
14 - Paths to the training and validation datasets,
15 - The number of classes
16 - The list of class names.
17 """
19 train: str = Field(..., description="Path to the training dataset")
20 val: str = Field(..., description="Path to the validation dataset")
21 nc: int = Field(..., description="Number of classes")
22 names: list[str] = Field(..., description="List of class names")
24 @classmethod
25 def from_yaml(cls, yaml_path: Path) -> "YOLOConfig":
26 """Load YOLO configuration from a YAML file.
28 Args:
29 yaml_path (Path): Path to the YAML configuration file.
31 Returns:
32 YOLOConfig: The loaded YOLO configuration.
33 """
34 import yaml
36 with open(yaml_path, "r") as f:
37 data = yaml.safe_load(f)
38 return cls(**data)
41class YOLOConverter(Converter):
42 """
43 Converter for YOLO dataset format to Collectra format.
44 This converter reads the YOLO configuration file, processes the training and validation datasets, and converts the annotations to Collectra format.
45 The converted files are saved in the specified output directory.
46 """
48 input: Path
49 output: Path
50 root_label: str
51 ext: str
52 force: bool
54 def __init__(
55 self,
56 input: str | Path,
57 root_label: str,
58 ext: str,
59 output: str | Path,
60 force: bool = False,
61 ):
62 """Initialize the YOLOConverter. Process the input configuration file or folder to extract necessary information for conversion.
64 Args:
65 input (str | Path): Path to the YOLO configuration file or folder.
66 root_label (str): Root label for converted data.
67 ext (str): File extension for converted collectra files.
68 output (str | Path): Output folder for converted files.
69 force (bool, optional): Whether to force overwrite existing files. Defaults to False.
71 """
72 self.input = Path(input)
73 self.root_label = root_label
74 self.ext = ext
75 self.output = Path(output)
76 self.force = force
77 self._get_config()
79 def _get_config(self) -> YOLOConfig:
80 """Get the YOLO configuration from the input file or directory.
82 Returns:
83 YOLOConfig: The YOLO configuration.
85 """
86 if not (self.input.is_file() or self.input.is_dir()):
87 raise ValueError(
88 f"Input path {self.input} is neither a file nor a directory."
89 )
90 if self.input.is_file():
91 if self.input.suffix != ".yaml":
92 raise ValueError(
93 f"Expected a YAML config file, but got {self.input.suffix}"
94 )
95 if self.input.is_dir():
96 config_files = list(self.input.glob("*.yaml"))
97 if not config_files:
98 raise ValueError(
99 f"No YAML config files found in directory {self.input}"
100 )
101 if len(config_files) > 1:
102 raise ValueError(
103 f"Multiple YAML config files found in directory {self.input}. Please ensure there is only one config file."
104 )
105 self.input = config_files[0]
106 return YOLOConfig.from_yaml(self.input)
108 def _convert_file(
109 self, file: Path, txt_file: Path, partition: str | None = None
110 ) -> None:
111 collectra_file = CollectraFile.from_file(
112 file=file,
113 label=self.root_label,
114 ext=self.ext,
115 output=self.output,
116 force=self.force,
117 partition=partition,
118 )
119 """ Read the YOLO annotation file and convert it to Collectra format.
120 The YOLO annotation file contains lines with the format:
121 <class_id> <x_center> <y_center> <width_relative> <height_relative>
122 Save the converted data in the Collectra file.
124 Args:
125 file (Path): Path to the image file.
126 txt_file (Path): Path to the YOLO annotation file.
128 """
129 with open(txt_file, "r") as f:
130 crops = [line.strip().split(" ") for line in f if line.strip()]
132 import uuid
134 for crop in crops:
135 class_id, x_center, y_center, width_relative, height_relative = crop
136 class_id = int(class_id)
137 class_name = self._get_config().names[class_id]
138 if class_name not in collectra_file.data:
139 collectra_file.data[class_name] = []
140 if not isinstance(collectra_file.data[class_name], list):
141 collectra_file.data[class_name] = [collectra_file.data[class_name]]
142 collectra_file.data[class_name].append(
143 {
144 "id": f"{class_name}-{uuid.uuid4().hex[:8]}",
145 "data": file.name,
146 "type": "collectra.ImageCrop",
147 "x_center": float(x_center),
148 "y_center": float(y_center),
149 "parents": self.root_label,
150 "width_relative": float(width_relative),
151 "height_relative": float(height_relative),
152 }
153 )
155 collectra_file.save()
157 def _process_files(self, input_of_files: Path, partition):
158 """Process the files listed in the input file.
159 The input file contains paths to image files, one per line.
160 For each image file, find the corresponding YOLO annotation file and convert it to Collectra format.
162 Args:
163 input_of_files (Path): Path to the file containing paths to image files.
165 Raises:
166 ValueError: If any of the image files or corresponding annotation files do not exist.
167 RuntimeError: If there is an error during file processing, the output directory will be cleaned
169 """
170 files: list[Path] = []
171 text_files: list[Path] = []
172 with open(input_of_files, "r") as f:
173 for line in f:
174 if not line.strip():
175 continue
176 file_path = (
177 Path(line.strip())
178 if str(line.strip()).startswith("/")
179 else input_of_files.parent / line.strip()
180 )
181 txt_path = Path(str(file_path).replace("images", "labels")).with_suffix(
182 ".txt"
183 )
184 if not file_path.exists() or not txt_path.exists():
185 raise ValueError(f"File {file_path} or {txt_path} does not exist.")
186 files.append(file_path)
187 text_files.append(txt_path)
188 try:
189 self.output.mkdir(parents=True, exist_ok=True)
190 for file, txt_file in zip(files, text_files):
191 self._convert_file(file, txt_file, partition=partition)
192 except Exception as e:
193 shutil.rmtree(self.output, ignore_errors=True)
194 raise RuntimeError(f"Failed to process files: {e}")
196 def convert(self) -> None:
197 """Convert the YOLO dataset to Collectra format. Process both training and validation datasets as specified in the YOLO configuration."""
198 config = self._get_config()
199 self._process_files(
200 self.input.parent / config.train, partition=str(Path(config.train).stem)
201 )
202 self._process_files(
203 self.input.parent / config.val, partition=str(Path(config.val).stem)
204 )
207class YOLOClassifierCSV:
208 """
209 Converter for YOLO classifier dataset format to Collectra format.
210 This converter extends the YOLOConverter to handle classification datasets in YOLO format.
211 It processes the training and validation datasets, converts the annotations to Collectra format, and saves the converted files in the specified output directory.
212 """
214 def __init__(
215 self,
216 input: str | Path,
217 root_label: str,
218 ext: str,
219 output: str | Path,
220 force: bool = False,
221 ):
222 """Initialize the YOLOConverter. Process the input configuration file or folder to extract necessary information for conversion.
224 Args:
225 input (str | Path): Path to the YOLO configuration file or folder.
226 root_label (str): Root label for converted data.
227 ext (str): File extension for converted collectra files.
228 output (str | Path): Output folder for converted files.
229 force (bool, optional): Whether to force overwrite existing files. Defaults to False.
231 """
232 self.input = Path(input)
233 self.root_label = root_label
234 self.ext = ext
235 self.output = Path(output)
236 self.force = force
238 def convert(self) -> None:
239 import pandas as pd
241 config_file = pd.read_csv(self.input)
242 for file_data in config_file.itertuples():
243 source_file = self.input.parent / str(file_data.path)
244 label_data = str(file_data.tag)
245 validation = "validation" if file_data.validation else "train"
246 print(
247 f"Processing file {source_file} with label {label_data} for partition {validation}"
248 )
249 self._convert_file(source_file, label_data, partition=validation)
251 def _convert_file(
252 self, file: Path, label_data: str, partition: str | None = None
253 ) -> None:
254 import uuid
256 collectra_file = CollectraFile.from_file(
257 file=file,
258 label=self.root_label,
259 ext=self.ext,
260 output=self.output,
261 force=self.force,
262 partition=partition,
263 )
264 collectra_file.data[f"{self.root_label}_writing_type"] = {
265 "id": f"{self.root_label}_writing_type-{uuid.uuid4().hex[:8]}",
266 "type": "collectra.Text",
267 "parents": collectra_file.data[self.root_label]["id"],
268 "data": label_data,
269 }
270 collectra_file.save()