Skip to content

Commit

Permalink
Apply linting suggestions for convert to meshes task
Browse files Browse the repository at this point in the history
  • Loading branch information
@jessicasyu
jessicasyu committed Sep 17, 2024
1 parent 3d5c792 commit e0d77d5
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 48 deletions.
87 changes: 55 additions & 32 deletions src/arcade_collection/convert/convert_to_meshes.py
View file
Original file line number Diff line number Diff line change
@@ -1,26 +1,42 @@
import tarfile
from typing import Optional, Union
from __future__ import annotations

from enum import Enum
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd
from skimage import measure

from arcade_collection.output.extract_tick_json import extract_tick_json
from arcade_collection.output.get_location_voxels import get_location_voxels

if TYPE_CHECKING:
import tarfile

import pandas as pd

MAX_ARRAY_LEVEL = 7
"""Maximum array level for conversion to meshes."""


class MeshType(Enum):
"""Mesh face types."""

DEFAULT = False
"""Mesh with default faces."""

INVERTED = True
"""Mesh with inverted faces."""


def convert_to_meshes(
series_key: str,
locations_tar: tarfile.TarFile,
frame_spec: tuple[int, int, int],
regions: list[str],
box: tuple[int, int, int],
invert: Union[bool, dict[str, bool]] = False,
group_size: Optional[int] = None,
categories: Optional[pd.DataFrame] = None,
mesh_type: MeshType | dict[str, MeshType] = MeshType.DEFAULT,
group_size: int | None = None,
categories: pd.DataFrame | None = None,
) -> list[tuple[int, int, str, str]]:
"""
Convert data to mesh OBJ contents.
Expand All @@ -37,8 +53,8 @@ def convert_to_meshes(
List of regions.
box
Size of bounding box.
invert
True to invert the order of faces, False otherwise.
mesh_type
Mesh face type.
group_size
Number of objects in each group (if grouping meshes).
categories
Expand All @@ -54,23 +70,19 @@ def convert_to_meshes(
meshes = []

length, width, height = box

if group_size is not None:
groups = make_mesh_groups(categories, frames, group_size)
else:
groups = None
groups = make_mesh_groups(categories, frames, group_size) if group_size is not None else None

for frame in frames:
locations = extract_tick_json(locations_tar, series_key, frame, "LOCATIONS")

for region in regions:
region_invert = invert[region] if isinstance(invert, dict) else invert
region_mesh_type = mesh_type[region] if isinstance(mesh_type, dict) else mesh_type

if groups is None:
for location in locations:
location_id = location["id"]
mesh = make_individual_mesh(
location, length, width, height, region, region_invert
location, length, width, height, region, region_mesh_type
)

if mesh is None:
Expand All @@ -83,7 +95,7 @@ def convert_to_meshes(
location for location in locations if location["id"] in group
]
mesh = make_combined_mesh(
group_locations, length, width, height, region, region_invert
group_locations, length, width, height, region, region_mesh_type
)

if mesh is None:
Expand Down Expand Up @@ -133,8 +145,13 @@ def make_mesh_groups(


def make_individual_mesh(
location: dict, length: int, width: int, height: int, region: str, invert: bool
) -> Optional[str]:
location: dict,
length: int,
width: int,
height: int,
region: str,
mesh_type: MeshType = MeshType.DEFAULT,
) -> str | None:
"""
Create mesh containing a single object.
Expand All @@ -150,8 +167,8 @@ def make_individual_mesh(
Bounding box height.
region
Region name.
invert
True to invert the order of faces, False otherwise.
mesh_type
Mesh face type.
Returns
-------
Expand All @@ -170,14 +187,17 @@ def make_individual_mesh(
center = list(np.array(voxels).mean(axis=0))
array = make_mesh_array(voxels, length, width, height)
verts, faces, normals = make_mesh_geometry(array, center)
mesh = make_mesh_file(verts, faces, normals, invert)

return mesh
return make_mesh_file(verts, faces, normals, mesh_type)


def make_combined_mesh(
locations: list[dict], length: int, width: int, height: int, region: str, invert: bool
) -> Optional[str]:
locations: list[dict],
length: int,
width: int,
height: int,
region: str,
mesh_type: MeshType = MeshType.DEFAULT,
) -> str | None:
"""
Create mesh containing multiple objects.
Expand All @@ -193,8 +213,8 @@ def make_combined_mesh(
Bounding box height.
region
Region name.
invert
True to invert the order of faces, False otherwise.
mesh_type
Mesh face type.
Returns
-------
Expand All @@ -217,7 +237,7 @@ def make_combined_mesh(
center = [length / 2, width / 2, height / 2]
array = make_mesh_array(voxels, length, width, height)
verts, faces, normals = make_mesh_geometry(array, center, offset)
mesh = make_mesh_file(verts, faces, normals, invert)
mesh = make_mesh_file(verts, faces, normals, mesh_type)

meshes.append(mesh)
offset = offset + len(verts)
Expand Down Expand Up @@ -318,7 +338,10 @@ def make_mesh_geometry(


def make_mesh_file(
verts: np.ndarray, faces: np.ndarray, normals: np.ndarray, invert: bool = False
verts: np.ndarray,
faces: np.ndarray,
normals: np.ndarray,
mesh_type: MeshType = MeshType.DEFAULT,
) -> str:
"""
Create mesh OBJ file contents from marching cubes output.
Expand All @@ -333,8 +356,8 @@ def make_mesh_file(
Array of mesh faces.
normals
Array of mesh normals.
invert
True to invert the order of faces, False otherwise.
mesh_type
Mesh face type.
Returns
-------
Expand All @@ -351,7 +374,7 @@ def make_mesh_file(
mesh += f"vn {item[0]} {item[1]} {item[2]}\n"

for item in faces:
if invert:
if mesh_type == MeshType.INVERTED:
mesh += f"f {item[0]}//{item[0]} {item[1]}//{item[1]} {item[2]}//{item[2]}\n"
else:
mesh += f"f {item[2]}//{item[2]} {item[1]}//{item[1]} {item[0]}//{item[0]}\n"
Expand Down
32 changes: 16 additions & 16 deletions tests/arcade_collection/convert/test_convert_to_meshes.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import numpy as np
import pandas as pd

from arcade_collection.convert.convert_to_meshes import MAX_ARRAY_LEVEL, convert_to_meshes
from arcade_collection.convert.convert_to_meshes import MAX_ARRAY_LEVEL, MeshType, convert_to_meshes

from .utilities import build_tar_instance

Expand Down Expand Up @@ -73,8 +73,8 @@ def setUp(self):
"measure",
return_value=mock.Mock(),
)
def test_convert_to_meshes_no_invert_no_group(self, measure_mock):
invert = False
def test_convert_to_meshes_default_no_group(self, measure_mock):
mesh_type = MeshType.DEFAULT
group_size = None
categories = None

Expand Down Expand Up @@ -122,7 +122,7 @@ def test_convert_to_meshes_no_invert_no_group(self, measure_mock):
self.frame_spec,
self.regions,
self.box,
invert,
mesh_type,
group_size,
categories,
)
Expand All @@ -134,8 +134,8 @@ def test_convert_to_meshes_no_invert_no_group(self, measure_mock):
"measure",
return_value=mock.Mock(),
)
def test_convert_to_meshes_all_invert_no_group(self, measure_mock):
invert = True
def test_convert_to_meshes_all_mesh_type_no_group(self, measure_mock):
mesh_type = MeshType.INVERTED
group_size = None
categories = None

Expand Down Expand Up @@ -183,7 +183,7 @@ def test_convert_to_meshes_all_invert_no_group(self, measure_mock):
self.frame_spec,
self.regions,
self.box,
invert,
mesh_type,
group_size,
categories,
)
Expand All @@ -195,8 +195,8 @@ def test_convert_to_meshes_all_invert_no_group(self, measure_mock):
"measure",
return_value=mock.Mock(),
)
def test_convert_to_meshes_region_invert_no_group(self, measure_mock):
invert = {"DEFAULT": False, "REGION": True}
def test_convert_to_meshes_region_mesh_type_no_group(self, measure_mock):
mesh_type = {"DEFAULT": MeshType.DEFAULT, "REGION": MeshType.INVERTED}
group_size = None
categories = None

Expand Down Expand Up @@ -257,7 +257,7 @@ def test_convert_to_meshes_region_invert_no_group(self, measure_mock):
self.frame_spec,
self.regions,
self.box,
invert,
mesh_type,
group_size,
categories,
)
Expand All @@ -269,8 +269,8 @@ def test_convert_to_meshes_region_invert_no_group(self, measure_mock):
"measure",
return_value=mock.Mock(),
)
def test_convert_to_meshes_no_invert_with_group_same_category(self, measure_mock):
invert = False
def test_convert_to_meshes_default_mesh_type_with_group_same_category(self, measure_mock):
mesh_type = MeshType.DEFAULT
group_size = 2
categories = pd.DataFrame(
{"FRAME": [5, 5, 10, 15], "CATEGORY": ["A", "A", "A", "A"], "ID": [1, 2, 3, 4]}
Expand Down Expand Up @@ -344,7 +344,7 @@ def test_convert_to_meshes_no_invert_with_group_same_category(self, measure_mock
self.frame_spec,
self.regions,
self.box,
invert,
mesh_type,
group_size,
categories,
)
Expand All @@ -356,8 +356,8 @@ def test_convert_to_meshes_no_invert_with_group_same_category(self, measure_mock
"measure",
return_value=mock.Mock(),
)
def test_convert_to_meshes_no_invert_with_group_different_category(self, measure_mock):
invert = False
def test_convert_to_meshes_default_mesh_type_with_group_different_category(self, measure_mock):
mesh_type = MeshType.DEFAULT
group_size = 2
categories = pd.DataFrame(
{"FRAME": [5, 5, 10, 15], "CATEGORY": ["A", "B", "A", "A"], "ID": [1, 2, 3, 4]}
Expand Down Expand Up @@ -420,7 +420,7 @@ def test_convert_to_meshes_no_invert_with_group_different_category(self, measure
self.frame_spec,
self.regions,
self.box,
invert,
mesh_type,
group_size,
categories,
)
Expand Down

0 comments on commit e0d77d5

Please sign in to comment.