hafnia 0.4.2__py3-none-any.whl → 0.4.3__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- hafnia/dataset/{dataset_upload_helper.py → dataset_details_uploader.py} +114 -191
- hafnia/dataset/dataset_names.py +26 -0
- hafnia/dataset/format_conversions/format_coco.py +490 -0
- hafnia/dataset/format_conversions/format_helpers.py +33 -0
- hafnia/dataset/format_conversions/format_image_classification_folder.py +95 -14
- hafnia/dataset/format_conversions/format_yolo.py +115 -25
- hafnia/dataset/format_conversions/torchvision_datasets.py +10 -8
- hafnia/dataset/hafnia_dataset.py +20 -466
- hafnia/dataset/hafnia_dataset_types.py +477 -0
- hafnia/dataset/license_types.py +4 -4
- hafnia/dataset/operations/dataset_stats.py +3 -3
- hafnia/dataset/operations/dataset_transformations.py +14 -17
- hafnia/dataset/operations/table_transformations.py +20 -13
- hafnia/dataset/primitives/bbox.py +6 -2
- hafnia/dataset/primitives/bitmask.py +21 -46
- hafnia/dataset/primitives/classification.py +1 -1
- hafnia/dataset/primitives/polygon.py +43 -2
- hafnia/dataset/primitives/primitive.py +1 -1
- hafnia/dataset/primitives/segmentation.py +1 -1
- hafnia/experiment/hafnia_logger.py +13 -4
- hafnia/platform/datasets.py +2 -3
- hafnia/torch_helpers.py +48 -4
- hafnia/utils.py +34 -0
- hafnia/visualizations/image_visualizations.py +3 -1
- {hafnia-0.4.2.dist-info → hafnia-0.4.3.dist-info}/METADATA +2 -2
- {hafnia-0.4.2.dist-info → hafnia-0.4.3.dist-info}/RECORD +29 -26
- {hafnia-0.4.2.dist-info → hafnia-0.4.3.dist-info}/WHEEL +0 -0
- {hafnia-0.4.2.dist-info → hafnia-0.4.3.dist-info}/entry_points.txt +0 -0
- {hafnia-0.4.2.dist-info → hafnia-0.4.3.dist-info}/licenses/LICENSE +0 -0
|
@@ -1,8 +1,7 @@
|
|
|
1
1
|
from pathlib import Path
|
|
2
|
-
from typing import
|
|
2
|
+
from typing import List, Optional, Tuple, Type
|
|
3
3
|
|
|
4
4
|
import polars as pl
|
|
5
|
-
from rich.progress import track
|
|
6
5
|
|
|
7
6
|
from hafnia.dataset.dataset_names import (
|
|
8
7
|
FILENAME_ANNOTATIONS_JSONL,
|
|
@@ -10,14 +9,13 @@ from hafnia.dataset.dataset_names import (
|
|
|
10
9
|
PrimitiveField,
|
|
11
10
|
SampleField,
|
|
12
11
|
)
|
|
12
|
+
from hafnia.dataset.hafnia_dataset_types import TaskInfo
|
|
13
13
|
from hafnia.dataset.operations import table_transformations
|
|
14
14
|
from hafnia.dataset.primitives import PRIMITIVE_TYPES
|
|
15
15
|
from hafnia.dataset.primitives.classification import Classification
|
|
16
16
|
from hafnia.dataset.primitives.primitive import Primitive
|
|
17
17
|
from hafnia.log import user_logger
|
|
18
|
-
|
|
19
|
-
if TYPE_CHECKING:
|
|
20
|
-
from hafnia.dataset.hafnia_dataset import TaskInfo
|
|
18
|
+
from hafnia.utils import progress_bar
|
|
21
19
|
|
|
22
20
|
|
|
23
21
|
def create_primitive_table(
|
|
@@ -29,13 +27,11 @@ def create_primitive_table(
|
|
|
29
27
|
"""
|
|
30
28
|
Returns a DataFrame with objects of the specified primitive type.
|
|
31
29
|
"""
|
|
32
|
-
|
|
33
|
-
has_primitive_column = (column_name in samples_table.columns) and (
|
|
34
|
-
samples_table[column_name].dtype == pl.List(pl.Struct)
|
|
35
|
-
)
|
|
36
|
-
if not has_primitive_column:
|
|
30
|
+
if not has_primitive(samples_table, PrimitiveType):
|
|
37
31
|
return None
|
|
38
32
|
|
|
33
|
+
column_name = PrimitiveType.column_name()
|
|
34
|
+
|
|
39
35
|
# Remove frames without objects
|
|
40
36
|
remove_no_object_frames = samples_table.filter(pl.col(column_name).list.len() > 0)
|
|
41
37
|
|
|
@@ -60,6 +56,17 @@ def create_primitive_table(
|
|
|
60
56
|
return objects_df
|
|
61
57
|
|
|
62
58
|
|
|
59
|
+
def has_primitive(samples: pl.DataFrame, PrimitiveType: Type[Primitive]) -> bool:
|
|
60
|
+
col_name = PrimitiveType.column_name()
|
|
61
|
+
if col_name not in samples.columns:
|
|
62
|
+
return False
|
|
63
|
+
|
|
64
|
+
if samples[col_name].dtype != pl.List(pl.Struct):
|
|
65
|
+
return False
|
|
66
|
+
|
|
67
|
+
return True
|
|
68
|
+
|
|
69
|
+
|
|
63
70
|
def merge_samples(samples0: pl.DataFrame, samples1: pl.DataFrame) -> pl.DataFrame:
|
|
64
71
|
has_same_schema = samples0.schema == samples1.schema
|
|
65
72
|
if not has_same_schema:
|
|
@@ -215,7 +222,7 @@ def read_samples_from_path(path: Path) -> pl.DataFrame:
|
|
|
215
222
|
def check_image_paths(table: pl.DataFrame) -> bool:
|
|
216
223
|
missing_files = []
|
|
217
224
|
org_paths = table[SampleField.FILE_PATH].to_list()
|
|
218
|
-
for org_path in
|
|
225
|
+
for org_path in progress_bar(org_paths, description="Check image paths"):
|
|
219
226
|
org_path = Path(org_path)
|
|
220
227
|
if not org_path.exists():
|
|
221
228
|
missing_files.append(org_path)
|
|
@@ -273,7 +280,7 @@ def unnest_classification_tasks(table: pl.DataFrame, strict: bool = True) -> pl.
|
|
|
273
280
|
return table_out
|
|
274
281
|
|
|
275
282
|
|
|
276
|
-
def update_class_indices(samples: pl.DataFrame, task:
|
|
283
|
+
def update_class_indices(samples: pl.DataFrame, task: TaskInfo) -> pl.DataFrame:
|
|
277
284
|
if task.class_names is None or len(task.class_names) == 0:
|
|
278
285
|
raise ValueError(f"Task '{task.name}' does not have defined class names to update class indices.")
|
|
279
286
|
|
|
@@ -318,7 +325,7 @@ def add_sample_index(samples: pl.DataFrame) -> pl.DataFrame:
|
|
|
318
325
|
if SampleField.SAMPLE_INDEX in samples.columns:
|
|
319
326
|
samples = samples.drop(SampleField.SAMPLE_INDEX)
|
|
320
327
|
samples = samples.select(
|
|
321
|
-
pl.int_range(0, pl.
|
|
328
|
+
pl.int_range(0, pl.len(), dtype=pl.UInt64).alias(SampleField.SAMPLE_INDEX),
|
|
322
329
|
pl.all(),
|
|
323
330
|
)
|
|
324
331
|
return samples
|
|
@@ -30,6 +30,9 @@ class Bbox(Primitive):
|
|
|
30
30
|
top_left_y: float = Field(
|
|
31
31
|
description="Normalized y-coordinate of top-left corner (0.0=top edge, 1.0=bottom edge) as a fraction of image height"
|
|
32
32
|
)
|
|
33
|
+
area: Optional[float] = Field(
|
|
34
|
+
default=None, description="Area of the bounding box as a fraction of the image area (0.0 to 1.0)"
|
|
35
|
+
)
|
|
33
36
|
class_name: Optional[str] = Field(default=None, description="Class name, e.g. 'car'")
|
|
34
37
|
class_idx: Optional[int] = Field(default=None, description="Class index, e.g. 0 for 'car' if it is the first class")
|
|
35
38
|
object_id: Optional[str] = Field(default=None, description="Unique identifier for the object, e.g. '12345123'")
|
|
@@ -49,7 +52,8 @@ class Bbox(Primitive):
|
|
|
49
52
|
def column_name() -> str:
|
|
50
53
|
return "bboxes"
|
|
51
54
|
|
|
52
|
-
def calculate_area(self) -> float:
|
|
55
|
+
def calculate_area(self, image_height: int, image_width: int) -> float:
|
|
56
|
+
"""Calculates the area of the bounding box as a fraction of the image area."""
|
|
53
57
|
return self.height * self.width
|
|
54
58
|
|
|
55
59
|
@staticmethod
|
|
@@ -73,7 +77,7 @@ class Bbox(Primitive):
|
|
|
73
77
|
"""
|
|
74
78
|
return (self.top_left_x, self.top_left_y, self.width, self.height)
|
|
75
79
|
|
|
76
|
-
def
|
|
80
|
+
def to_coco_ints(self, image_height: int, image_width: int) -> Tuple[int, int, int, int]:
|
|
77
81
|
xmin = round_int_clip_value(self.top_left_x * image_width, max_value=image_width)
|
|
78
82
|
bbox_width = round_int_clip_value(self.width * image_width, max_value=image_width)
|
|
79
83
|
|
|
@@ -22,11 +22,11 @@ class Bitmask(Primitive):
|
|
|
22
22
|
left: int = Field(description="Bitmask left coordinate in pixels")
|
|
23
23
|
height: int = Field(description="Bitmask height of the bounding box in pixels")
|
|
24
24
|
width: int = Field(description="Bitmask width of the bounding box in pixels")
|
|
25
|
-
|
|
25
|
+
rle_string: str = Field(
|
|
26
26
|
description="Run-length encoding (RLE) string for the bitmask region of size (height, width) at (top, left)."
|
|
27
27
|
)
|
|
28
28
|
area: Optional[float] = Field(
|
|
29
|
-
default=None, description="Area of the bitmask
|
|
29
|
+
default=None, description="Area of the bitmask as a fraction of the image area (0.0 to 1.0)"
|
|
30
30
|
)
|
|
31
31
|
class_name: Optional[str] = Field(default=None, description="Class name of the object represented by the bitmask")
|
|
32
32
|
class_idx: Optional[int] = Field(default=None, description="Class index of the object represented by the bitmask")
|
|
@@ -47,8 +47,9 @@ class Bitmask(Primitive):
|
|
|
47
47
|
def column_name() -> str:
|
|
48
48
|
return "bitmasks"
|
|
49
49
|
|
|
50
|
-
def calculate_area(self) -> float:
|
|
51
|
-
|
|
50
|
+
def calculate_area(self, image_height: int, image_width: int) -> float:
|
|
51
|
+
area_px = coco_mask.area(self.to_coco_rle(img_height=image_height, img_width=image_width))
|
|
52
|
+
return area_px / (image_height * image_width)
|
|
52
53
|
|
|
53
54
|
@staticmethod
|
|
54
55
|
def from_mask(
|
|
@@ -79,60 +80,34 @@ class Bitmask(Primitive):
|
|
|
79
80
|
height=h,
|
|
80
81
|
width=w,
|
|
81
82
|
area=area,
|
|
82
|
-
|
|
83
|
+
rle_string=rle_string,
|
|
83
84
|
class_name=class_name,
|
|
84
85
|
class_idx=class_idx,
|
|
85
86
|
object_id=object_id,
|
|
86
87
|
)
|
|
87
88
|
|
|
88
|
-
def squeeze_mask(self):
|
|
89
|
-
"""
|
|
90
|
-
A mask may have large redundant areas of zeros. This function squeezes the mask to remove those areas.
|
|
91
|
-
"""
|
|
92
|
-
region_mask = self.to_region_mask()
|
|
93
|
-
shift_left, last_left = np.flatnonzero(region_mask.sum(axis=0))[[0, -1]]
|
|
94
|
-
shift_top, last_top = np.flatnonzero(region_mask.sum(axis=1))[[0, -1]]
|
|
95
|
-
new_top = self.top + shift_top
|
|
96
|
-
new_left = self.left + shift_left
|
|
97
|
-
new_region_mask = region_mask[shift_top : last_top + 1, shift_left : last_left + 1]
|
|
98
|
-
|
|
99
|
-
bitmask_squeezed = Bitmask.from_mask(
|
|
100
|
-
mask=new_region_mask,
|
|
101
|
-
top=new_top,
|
|
102
|
-
left=new_left,
|
|
103
|
-
class_name=self.class_name,
|
|
104
|
-
class_idx=self.class_idx,
|
|
105
|
-
object_id=self.object_id,
|
|
106
|
-
)
|
|
107
|
-
return bitmask_squeezed
|
|
108
|
-
|
|
109
89
|
def anonymize_by_blurring(self, image: np.ndarray, inplace: bool = False, max_resolution: int = 20) -> np.ndarray:
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
region_image = image[
|
|
114
|
-
mask_tight.top : mask_tight.top + mask_tight.height, mask_tight.left : mask_tight.left + mask_tight.width
|
|
115
|
-
]
|
|
90
|
+
mask = self.to_mask(img_height=image.shape[0], img_width=image.shape[1])
|
|
91
|
+
region_mask = mask[self.top : self.top + self.height, self.left : self.left + self.width]
|
|
92
|
+
region_image = image[self.top : self.top + self.height, self.left : self.left + self.width]
|
|
116
93
|
region_image_blurred = anonymize_by_resizing(blur_region=region_image, max_resolution=max_resolution)
|
|
117
|
-
image_mixed = np.where(
|
|
118
|
-
image[
|
|
119
|
-
mask_tight.top : mask_tight.top + mask_tight.height, mask_tight.left : mask_tight.left + mask_tight.width
|
|
120
|
-
] = image_mixed
|
|
94
|
+
image_mixed = np.where(region_mask[:, :, None], region_image_blurred, region_image)
|
|
95
|
+
image[self.top : self.top + self.height, self.left : self.left + self.width] = image_mixed
|
|
121
96
|
return image
|
|
122
97
|
|
|
123
|
-
def
|
|
124
|
-
"""Returns
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
98
|
+
def to_coco_rle(self, img_height: int, img_width: int, as_bytes: bool = True) -> Dict[str, Any]:
|
|
99
|
+
"""Returns the COCO RLE dictionary from the RLE string."""
|
|
100
|
+
|
|
101
|
+
rle_string = self.rle_string
|
|
102
|
+
if as_bytes:
|
|
103
|
+
rle_string = rle_string.encode() # type: ignore[assignment]
|
|
104
|
+
rle = {"counts": rle_string, "size": [img_height, img_width]}
|
|
105
|
+
return rle
|
|
128
106
|
|
|
129
107
|
def to_mask(self, img_height: int, img_width: int) -> np.ndarray:
|
|
130
108
|
"""Creates a full image mask from the RLE string."""
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
bitmask_np = np.zeros((img_height, img_width), dtype=bool)
|
|
134
|
-
bitmask_np[self.top : self.top + self.height, self.left : self.left + self.width] = region_mask
|
|
135
|
-
return bitmask_np
|
|
109
|
+
mask = coco_mask.decode(self.to_coco_rle(img_height=img_height, img_width=img_width)) > 0
|
|
110
|
+
return mask
|
|
136
111
|
|
|
137
112
|
def draw(self, image: np.ndarray, inplace: bool = False, draw_label: bool = True) -> np.ndarray:
|
|
138
113
|
if not inplace:
|
|
@@ -33,7 +33,7 @@ class Classification(Primitive):
|
|
|
33
33
|
def column_name() -> str:
|
|
34
34
|
return "classifications"
|
|
35
35
|
|
|
36
|
-
def calculate_area(self) -> float:
|
|
36
|
+
def calculate_area(self, image_height: int, image_width: int) -> float:
|
|
37
37
|
return 1.0
|
|
38
38
|
|
|
39
39
|
def draw(self, image: np.ndarray, inplace: bool = False, draw_label: bool = True) -> np.ndarray:
|
|
@@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple
|
|
|
2
2
|
|
|
3
3
|
import cv2
|
|
4
4
|
import numpy as np
|
|
5
|
+
from more_itertools import collapse
|
|
6
|
+
from pycocotools import mask as coco_utils
|
|
5
7
|
from pydantic import Field
|
|
6
8
|
|
|
7
9
|
from hafnia.dataset.primitives.bitmask import Bitmask
|
|
@@ -13,6 +15,7 @@ from hafnia.dataset.primitives.utils import class_color_by_name, get_class_name
|
|
|
13
15
|
class Polygon(Primitive):
|
|
14
16
|
# Names should match names in FieldName
|
|
15
17
|
points: List[Point] = Field(description="List of points defining the polygon")
|
|
18
|
+
area: Optional[float] = Field(default=None, description="Area of the polygon in pixels")
|
|
16
19
|
class_name: Optional[str] = Field(default=None, description="Class name of the polygon")
|
|
17
20
|
class_idx: Optional[int] = Field(default=None, description="Class index of the polygon")
|
|
18
21
|
object_id: Optional[str] = Field(default=None, description="Object ID of the polygon")
|
|
@@ -44,7 +47,7 @@ class Polygon(Primitive):
|
|
|
44
47
|
def column_name() -> str:
|
|
45
48
|
return "polygons"
|
|
46
49
|
|
|
47
|
-
def calculate_area(self) -> float:
|
|
50
|
+
def calculate_area(self, image_height: int, image_width: int) -> float:
|
|
48
51
|
raise NotImplementedError()
|
|
49
52
|
|
|
50
53
|
def to_pixel_coordinates(
|
|
@@ -81,11 +84,49 @@ class Polygon(Primitive):
|
|
|
81
84
|
points = np.array(self.to_pixel_coordinates(image_shape=image.shape[:2]))
|
|
82
85
|
mask = np.zeros(image.shape[:2], dtype=np.uint8)
|
|
83
86
|
mask = cv2.fillPoly(mask, [points], color=255).astype(bool)
|
|
84
|
-
bitmask = Bitmask.from_mask(mask=mask, top=0, left=0)
|
|
87
|
+
bitmask = Bitmask.from_mask(mask=mask, top=0, left=0)
|
|
85
88
|
image = bitmask.anonymize_by_blurring(image=image, inplace=inplace, max_resolution=max_resolution)
|
|
86
89
|
|
|
87
90
|
return image
|
|
88
91
|
|
|
92
|
+
def to_mask(self, img_height: int, img_width: int, use_coco_utils=False) -> np.ndarray:
|
|
93
|
+
if use_coco_utils:
|
|
94
|
+
points = list(collapse(self.to_pixel_coordinates(image_shape=(img_height, img_width))))
|
|
95
|
+
rles = coco_utils.frPyObjects([points], img_height, img_width)
|
|
96
|
+
rle = coco_utils.merge(rles)
|
|
97
|
+
mask = coco_utils.decode(rle).astype(bool)
|
|
98
|
+
return mask
|
|
99
|
+
|
|
100
|
+
mask = np.zeros((img_height, img_width), dtype=np.uint8)
|
|
101
|
+
points = np.array(self.to_pixel_coordinates(image_shape=(img_height, img_width)))
|
|
102
|
+
mask = cv2.fillPoly(mask, [points], color=255).astype(bool)
|
|
103
|
+
return mask
|
|
104
|
+
|
|
105
|
+
def to_bitmask(self, img_height: int, img_width: int) -> Bitmask:
|
|
106
|
+
points = list(collapse(self.to_pixel_coordinates(image_shape=(img_height, img_width))))
|
|
107
|
+
rles = coco_utils.frPyObjects([points], img_height, img_width)
|
|
108
|
+
rle = coco_utils.merge(rles)
|
|
109
|
+
top, left, height, width = coco_utils.toBbox(rle)
|
|
110
|
+
|
|
111
|
+
rle_string = rle["counts"]
|
|
112
|
+
if isinstance(rle_string, bytes):
|
|
113
|
+
rle_string = rle_string.decode("utf-8")
|
|
114
|
+
|
|
115
|
+
return Bitmask(
|
|
116
|
+
rle_string=rle_string,
|
|
117
|
+
top=int(top),
|
|
118
|
+
left=int(left),
|
|
119
|
+
width=int(width),
|
|
120
|
+
height=int(height),
|
|
121
|
+
class_name=self.class_name,
|
|
122
|
+
class_idx=self.class_idx,
|
|
123
|
+
object_id=self.object_id,
|
|
124
|
+
confidence=self.confidence,
|
|
125
|
+
ground_truth=self.ground_truth,
|
|
126
|
+
task_name=self.task_name,
|
|
127
|
+
meta=self.meta,
|
|
128
|
+
)
|
|
129
|
+
|
|
89
130
|
def mask(
|
|
90
131
|
self, image: np.ndarray, inplace: bool = False, color: Optional[Tuple[np.uint8, np.uint8, np.uint8]] = None
|
|
91
132
|
) -> np.ndarray:
|
|
@@ -30,7 +30,7 @@ class Segmentation(Primitive):
|
|
|
30
30
|
def column_name() -> str:
|
|
31
31
|
return "segmentations"
|
|
32
32
|
|
|
33
|
-
def calculate_area(self) -> float:
|
|
33
|
+
def calculate_area(self, image_height: int, image_width: int) -> float:
|
|
34
34
|
raise NotImplementedError()
|
|
35
35
|
|
|
36
36
|
def draw(self, image: np.ndarray, inplace: bool = False) -> np.ndarray:
|
|
@@ -85,8 +85,9 @@ class Entity(BaseModel):
|
|
|
85
85
|
class HafniaLogger:
|
|
86
86
|
EXPERIMENT_FILE = "experiment.parquet"
|
|
87
87
|
|
|
88
|
-
def __init__(self, log_dir: Union[Path, str] = "./.data"):
|
|
88
|
+
def __init__(self, project_name: str, log_dir: Union[Path, str] = "./.data"):
|
|
89
89
|
self._local_experiment_path = Path(log_dir) / "experiments" / now_as_str()
|
|
90
|
+
self.project_name = project_name
|
|
90
91
|
create_paths = [
|
|
91
92
|
self._local_experiment_path,
|
|
92
93
|
self.path_model_checkpoints(),
|
|
@@ -109,6 +110,7 @@ class HafniaLogger:
|
|
|
109
110
|
self._init_mlflow()
|
|
110
111
|
|
|
111
112
|
self.log_environment()
|
|
113
|
+
self.log_configuration({"project_name": project_name})
|
|
112
114
|
|
|
113
115
|
def _init_mlflow(self):
|
|
114
116
|
"""Initialize MLflow tracking for remote jobs."""
|
|
@@ -125,9 +127,16 @@ class HafniaLogger:
|
|
|
125
127
|
mlflow.set_experiment(experiment_name)
|
|
126
128
|
user_logger.info(f"MLflow experiment set to: {experiment_name}")
|
|
127
129
|
|
|
128
|
-
# Start MLflow run
|
|
130
|
+
# Start MLflow run with tags
|
|
129
131
|
run_name = os.getenv("MLFLOW_RUN_NAME", "undefined")
|
|
130
|
-
|
|
132
|
+
created_by = os.getenv("MLFLOW_CREATED_BY")
|
|
133
|
+
tags = {"project_name": self.project_name}
|
|
134
|
+
if experiment_name:
|
|
135
|
+
tags["organization_id"] = experiment_name
|
|
136
|
+
if created_by:
|
|
137
|
+
tags["created_by"] = created_by
|
|
138
|
+
|
|
139
|
+
mlflow.start_run(run_name=run_name, tags=tags, log_system_metrics=True)
|
|
131
140
|
self._mlflow_initialized = True
|
|
132
141
|
user_logger.info("MLflow run started successfully")
|
|
133
142
|
|
|
@@ -290,7 +299,7 @@ def get_instructions_how_to_store_model() -> str:
|
|
|
290
299
|
from hafnia.experiment import HafniaLogger
|
|
291
300
|
|
|
292
301
|
# Initiate Hafnia logger
|
|
293
|
-
logger = HafniaLogger()
|
|
302
|
+
logger = HafniaLogger(project_name="my_classification_project")
|
|
294
303
|
|
|
295
304
|
# Folder path to store models - generated by the hafnia logger.
|
|
296
305
|
model_dir = logger.path_model()
|
hafnia/platform/datasets.py
CHANGED
|
@@ -9,7 +9,6 @@ from typing import Any, Dict, List, Optional
|
|
|
9
9
|
|
|
10
10
|
import rich
|
|
11
11
|
from rich import print as rprint
|
|
12
|
-
from rich.progress import track
|
|
13
12
|
|
|
14
13
|
from hafnia import http, utils
|
|
15
14
|
from hafnia.dataset.dataset_names import DATASET_FILENAMES_REQUIRED
|
|
@@ -21,7 +20,7 @@ from hafnia.dataset.hafnia_dataset import HafniaDataset
|
|
|
21
20
|
from hafnia.http import fetch
|
|
22
21
|
from hafnia.log import sys_logger, user_logger
|
|
23
22
|
from hafnia.platform.download import get_resource_credentials
|
|
24
|
-
from hafnia.utils import timed
|
|
23
|
+
from hafnia.utils import progress_bar, timed
|
|
25
24
|
from hafnia_cli.config import Config
|
|
26
25
|
|
|
27
26
|
|
|
@@ -192,7 +191,7 @@ def execute_s5cmd_commands(
|
|
|
192
191
|
|
|
193
192
|
error_lines = []
|
|
194
193
|
lines = []
|
|
195
|
-
for line in
|
|
194
|
+
for line in progress_bar(process.stdout, total=len(commands), description=description): # type: ignore[arg-type]
|
|
196
195
|
if "ERROR" in line or "error" in line:
|
|
197
196
|
error_lines.append(line.strip())
|
|
198
197
|
lines.append(line.strip())
|
hafnia/torch_helpers.py
CHANGED
|
@@ -1,6 +1,8 @@
|
|
|
1
1
|
from typing import Dict, List, Optional, Tuple, Type, Union
|
|
2
2
|
|
|
3
|
+
import cv2
|
|
3
4
|
import numpy as np
|
|
5
|
+
import polars as pl
|
|
4
6
|
import torch
|
|
5
7
|
import torchvision
|
|
6
8
|
from flatten_dict import flatten, unflatten
|
|
@@ -9,8 +11,9 @@ from torchvision import tv_tensors
|
|
|
9
11
|
from torchvision import utils as tv_utils
|
|
10
12
|
from torchvision.transforms import v2
|
|
11
13
|
|
|
12
|
-
from hafnia.dataset.dataset_names import PrimitiveField
|
|
13
|
-
from hafnia.dataset.hafnia_dataset import HafniaDataset
|
|
14
|
+
from hafnia.dataset.dataset_names import PrimitiveField, SampleField
|
|
15
|
+
from hafnia.dataset.hafnia_dataset import HafniaDataset
|
|
16
|
+
from hafnia.dataset.hafnia_dataset_types import Sample
|
|
14
17
|
from hafnia.dataset.primitives import (
|
|
15
18
|
PRIMITIVE_COLUMN_NAMES,
|
|
16
19
|
class_color_by_name,
|
|
@@ -18,6 +21,7 @@ from hafnia.dataset.primitives import (
|
|
|
18
21
|
from hafnia.dataset.primitives.bbox import Bbox
|
|
19
22
|
from hafnia.dataset.primitives.bitmask import Bitmask
|
|
20
23
|
from hafnia.dataset.primitives.classification import Classification
|
|
24
|
+
from hafnia.dataset.primitives.polygon import Polygon
|
|
21
25
|
from hafnia.dataset.primitives.primitive import Primitive
|
|
22
26
|
from hafnia.dataset.primitives.segmentation import Segmentation
|
|
23
27
|
from hafnia.log import user_logger
|
|
@@ -50,6 +54,16 @@ class TorchvisionDataset(torch.utils.data.Dataset):
|
|
|
50
54
|
):
|
|
51
55
|
self.dataset = dataset
|
|
52
56
|
|
|
57
|
+
self.max_points_in_polygon = 0
|
|
58
|
+
|
|
59
|
+
if self.dataset.has_primitive(Polygon):
|
|
60
|
+
self.max_points_in_polygon = (
|
|
61
|
+
self.dataset.samples[SampleField.POLYGONS]
|
|
62
|
+
.list.eval(pl.element().struct.field("points").list.len())
|
|
63
|
+
.explode()
|
|
64
|
+
.max()
|
|
65
|
+
)
|
|
66
|
+
|
|
53
67
|
self.transforms = transforms
|
|
54
68
|
self.keep_metadata = keep_metadata
|
|
55
69
|
|
|
@@ -74,7 +88,7 @@ class TorchvisionDataset(torch.utils.data.Dataset):
|
|
|
74
88
|
|
|
75
89
|
bbox_tasks: Dict[str, List[Bbox]] = get_primitives_per_task_name_for_primitive(sample, Bbox)
|
|
76
90
|
for task_name, bboxes in bbox_tasks.items():
|
|
77
|
-
bboxes_list = [bbox.
|
|
91
|
+
bboxes_list = [bbox.to_coco_ints(image_height=h, image_width=w) for bbox in bboxes]
|
|
78
92
|
bboxes_tensor = torch.as_tensor(bboxes_list).reshape(-1, 4)
|
|
79
93
|
target_flat[f"{Bbox.column_name()}.{task_name}"] = {
|
|
80
94
|
PrimitiveField.CLASS_IDX: [bbox.class_idx for bbox in bboxes],
|
|
@@ -91,6 +105,22 @@ class TorchvisionDataset(torch.utils.data.Dataset):
|
|
|
91
105
|
"mask": tv_tensors.Mask(bitmasks_np),
|
|
92
106
|
}
|
|
93
107
|
|
|
108
|
+
polygon_tasks: Dict[str, List[Polygon]] = get_primitives_per_task_name_for_primitive(sample, Polygon)
|
|
109
|
+
for task_name, polygons in polygon_tasks.items():
|
|
110
|
+
polygon_tensors = [
|
|
111
|
+
torch.tensor(pg.to_pixel_coordinates(image_shape=(h, w), as_int=False)) for pg in polygons
|
|
112
|
+
]
|
|
113
|
+
n_polygons = len(polygons)
|
|
114
|
+
polygons_matrix = torch.full((n_polygons, self.max_points_in_polygon, 2), fill_value=torch.nan)
|
|
115
|
+
|
|
116
|
+
for i, polygon_tensor in enumerate(polygon_tensors):
|
|
117
|
+
polygons_matrix[i, : polygon_tensor.shape[0], :] = polygon_tensor
|
|
118
|
+
|
|
119
|
+
target_flat[f"{Polygon.column_name()}.{task_name}"] = {
|
|
120
|
+
PrimitiveField.CLASS_IDX: [polygon.class_idx for polygon in polygons],
|
|
121
|
+
PrimitiveField.CLASS_NAME: [polygon.class_name for polygon in polygons],
|
|
122
|
+
"polygon": tv_tensors.KeyPoints(polygons_matrix, canvas_size=(h, w)),
|
|
123
|
+
}
|
|
94
124
|
if self.transforms:
|
|
95
125
|
image, target_flat = self.transforms(image, target_flat)
|
|
96
126
|
|
|
@@ -181,6 +211,18 @@ def draw_image_and_targets(
|
|
|
181
211
|
colors=colors,
|
|
182
212
|
)
|
|
183
213
|
|
|
214
|
+
if Polygon.column_name() in targets:
|
|
215
|
+
primitive_annotations = targets[Polygon.column_name()]
|
|
216
|
+
np_image = visualize_image.permute(1, 2, 0).numpy()
|
|
217
|
+
for task_name, task_annotations in primitive_annotations.items():
|
|
218
|
+
task_annotations["polygon"]
|
|
219
|
+
colors = [class_color_by_name(class_name) for class_name in task_annotations[PrimitiveField.CLASS_NAME]]
|
|
220
|
+
for color, polygon in zip(colors, task_annotations["polygon"], strict=True):
|
|
221
|
+
single_polygon = np.array(polygon[~torch.isnan(polygon[:, 0]), :][None, :, :]).astype(int)
|
|
222
|
+
|
|
223
|
+
np_image = cv2.polylines(np_image, [single_polygon], isClosed=False, color=color, thickness=2)
|
|
224
|
+
visualize_image = torch.from_numpy(np_image).permute(2, 0, 1)
|
|
225
|
+
|
|
184
226
|
# Important that classification is drawn last as it will change image dimensions
|
|
185
227
|
if Classification.column_name() in targets:
|
|
186
228
|
primitive_annotations = targets[Classification.column_name()]
|
|
@@ -219,7 +261,7 @@ class TorchVisionCollateFn:
|
|
|
219
261
|
images, targets = tuple(zip(*batch, strict=False))
|
|
220
262
|
if "image" not in self.skip_stacking_list:
|
|
221
263
|
images = torch.stack(images)
|
|
222
|
-
|
|
264
|
+
height, width = images.shape[-2:]
|
|
223
265
|
keys_min = set(targets[0])
|
|
224
266
|
keys_max = set(targets[0])
|
|
225
267
|
for target in targets:
|
|
@@ -250,6 +292,8 @@ class TorchVisionCollateFn:
|
|
|
250
292
|
item_values = tv_tensors.Image(item_values)
|
|
251
293
|
elif isinstance(first_element, tv_tensors.BoundingBoxes):
|
|
252
294
|
item_values = tv_tensors.BoundingBoxes(item_values)
|
|
295
|
+
elif isinstance(first_element, tv_tensors.KeyPoints):
|
|
296
|
+
item_values = tv_tensors.KeyPoints(item_values, canvas_size=(height, width))
|
|
253
297
|
targets_modified[key_name] = item_values
|
|
254
298
|
|
|
255
299
|
return images, targets_modified
|
hafnia/utils.py
CHANGED
|
@@ -2,6 +2,7 @@ import hashlib
|
|
|
2
2
|
import os
|
|
3
3
|
import time
|
|
4
4
|
import zipfile
|
|
5
|
+
from collections.abc import Sized
|
|
5
6
|
from datetime import datetime
|
|
6
7
|
from functools import wraps
|
|
7
8
|
from pathlib import Path
|
|
@@ -13,6 +14,7 @@ import pathspec
|
|
|
13
14
|
import rich
|
|
14
15
|
import seedir
|
|
15
16
|
from rich import print as rprint
|
|
17
|
+
from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn, TimeElapsedColumn, TimeRemainingColumn
|
|
16
18
|
|
|
17
19
|
from hafnia.log import sys_logger, user_logger
|
|
18
20
|
|
|
@@ -222,3 +224,35 @@ def remove_duplicates_preserve_order(seq: Iterable) -> List:
|
|
|
222
224
|
def is_image_file(file_path: Path) -> bool:
|
|
223
225
|
image_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".gif")
|
|
224
226
|
return file_path.suffix.lower() in image_extensions
|
|
227
|
+
|
|
228
|
+
|
|
229
|
+
def progress_bar(sequence: Iterable, total: Optional[int] = None, description: str = "Working...") -> Iterable:
|
|
230
|
+
"""
|
|
231
|
+
Progress bar showing number of iterations being processed with ETA and elapsed time.
|
|
232
|
+
|
|
233
|
+
Example usage:
|
|
234
|
+
|
|
235
|
+
```python
|
|
236
|
+
items = list(range(1000))
|
|
237
|
+
for item in progress_bar(items, description="Processing..."):
|
|
238
|
+
time.sleep(0.02)
|
|
239
|
+
```
|
|
240
|
+
Processing... ━━━━━━━━━╸━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 245/1000 ETA: 0:00:16 | Elapsed: 0:00:05
|
|
241
|
+
"""
|
|
242
|
+
progress_bar = Progress(
|
|
243
|
+
TextColumn("{task.description}"),
|
|
244
|
+
BarColumn(),
|
|
245
|
+
MofNCompleteColumn(),
|
|
246
|
+
TextColumn("ETA:"),
|
|
247
|
+
TimeRemainingColumn(),
|
|
248
|
+
TextColumn("| Elapsed:"),
|
|
249
|
+
TimeElapsedColumn(),
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
if total is None:
|
|
253
|
+
total = len(sequence) if isinstance(sequence, Sized) else None
|
|
254
|
+
with progress_bar as progress:
|
|
255
|
+
task = progress.add_task(description, total=total)
|
|
256
|
+
for item in sequence:
|
|
257
|
+
yield item
|
|
258
|
+
progress.update(task, advance=1)
|
|
@@ -7,7 +7,7 @@ import numpy as np
|
|
|
7
7
|
import numpy.typing as npt
|
|
8
8
|
from PIL import Image
|
|
9
9
|
|
|
10
|
-
from hafnia.dataset.
|
|
10
|
+
from hafnia.dataset.hafnia_dataset_types import Sample
|
|
11
11
|
from hafnia.dataset.primitives import (
|
|
12
12
|
Bbox,
|
|
13
13
|
Bitmask,
|
|
@@ -175,6 +175,8 @@ def save_dataset_sample_set_visualizations(
|
|
|
175
175
|
draw_settings: Optional[Dict[Type[Primitive], Dict]] = None,
|
|
176
176
|
anonymize_settings: Optional[Dict[Type[Primitive], Dict]] = None,
|
|
177
177
|
) -> List[Path]:
|
|
178
|
+
from hafnia.dataset.hafnia_dataset import HafniaDataset
|
|
179
|
+
|
|
178
180
|
dataset = HafniaDataset.from_path(path_dataset)
|
|
179
181
|
shutil.rmtree(path_output_folder, ignore_errors=True)
|
|
180
182
|
path_output_folder.mkdir(parents=True)
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: hafnia
|
|
3
|
-
Version: 0.4.
|
|
3
|
+
Version: 0.4.3
|
|
4
4
|
Summary: Python SDK for communication with Hafnia platform.
|
|
5
5
|
Author-email: Milestone Systems <hafniaplatform@milestone.dk>
|
|
6
6
|
License-File: LICENSE
|
|
@@ -343,7 +343,7 @@ batch_size = 128
|
|
|
343
343
|
learning_rate = 0.001
|
|
344
344
|
|
|
345
345
|
# Initialize Hafnia logger
|
|
346
|
-
logger = HafniaLogger()
|
|
346
|
+
logger = HafniaLogger(project_name="my_classification_project")
|
|
347
347
|
|
|
348
348
|
# Log experiment parameters
|
|
349
349
|
logger.log_configuration({"batch_size": 128, "learning_rate": 0.001})
|