Skip to content

Commit

Permalink
as_image only show nonzero
Browse files Browse the repository at this point in the history
  • Loading branch information
Tumiz committed May 1, 2024
1 parent 0047800 commit 543380b
Showing 1 changed file with 19 additions and 13 deletions.
32 changes: 19 additions & 13 deletions py3d/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,18 @@ def read_npy(path) -> Vector:
return numpy.load(path).view(Vector)


def rand(*n) -> Vector | Vector2 | Vector3 | Vector4:
'''
Create a random vector with shape of n
'''
w = n[-1]
if w in [2, 3, 4]:
vtype = getattr(__module__, f"Vector{w}")
else:
vtype = Vector
return numpy.random.rand(*n).view(vtype)


class Vector(numpy.ndarray):
'''
Base class of Vector2, Vector3, Vector4 and Transform
Expand Down Expand Up @@ -287,11 +299,6 @@ def n(self):
else:
return self.shape

@classmethod
def rand(cls, *n) -> Vector | Vector2 | Vector3 | Vector4:
n += cls.BASE_SHAPE
return numpy.random.rand(*n).view(cls)

@property
def x(self) -> Vector:
return self[..., 0].view(Vector)
Expand Down Expand Up @@ -451,27 +458,26 @@ def to_csv(self, path):
def to_npy(self, path):
numpy.save(path, self)

def as_image(self, align_center=True, sample_rate=None):
def as_image(self, nonzero=True, align_center=True, sample_rate=None):
'''
Visualize the vector as an image, with mapped colors from black to yellow or the image's own colors
'''
if not sample_rate:
sample_rate = max(round(self.size / 1e6), 1)
sample = self[::-sample_rate, ::sample_rate]
if sample.dtype == numpy.uint8:
sample = sample / 255
sample = Color(sample / 255)
elif sample.ndim == 2:
sample = Color.map(sample)
c = numpy.ones(sample.shape[:-1] + (4,)).view(Color)
c[..., :min(sample.shape[-1], 4)] = sample
c = c.transpose(1, 0, 2)
w, h, _ = numpy.shape(c)
ret = Point(w, h)
*_, h, w = sample.n
ret = Point(*_, w, h)
if align_center:
ret.xyz = Vector3.grid(range(-w//2, w//2), range(-h//2, h//2))
else:
ret.xyz = Vector3.grid(range(w), range(h))
ret.color = c
ret.color = sample.transpose(1, 0, 2)
if nonzero:
return ret[ret.color.rgb.any(-1)]
return ret


Expand Down

0 comments on commit 543380b

Please sign in to comment.