Skip to content

Commit

Permalink
Satisfy mypy
Browse files Browse the repository at this point in the history
  • Loading branch information
druzsan committed May 10, 2024
1 parent e659cd2 commit 948f7c9
Show file tree
Hide file tree
Showing 4 changed files with 17 additions and 13 deletions.
1 change: 0 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,6 @@ plugins = "pydantic.mypy"
[[tool.mypy.overrides]]
module = [
"pygltflib",
"av",
"validators",
"h5py",
"pandas",
Expand Down
6 changes: 4 additions & 2 deletions renumics/spotlight/embeddings/preprocessors.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,10 @@
"""

import io
from typing import List
from typing import List, cast

import av
import av.audio
import numpy as np
import PIL.Image

Expand Down Expand Up @@ -36,10 +37,11 @@ def preprocess_audio_batch(
for raw_data in raw_values:
with av.open(io.BytesIO(raw_data), "r") as container:
resampler = av.AudioResampler(
format="dbl", layout="mono", rate=sampling_rate
format="dbl", layout="mono", rate=sampling_rate # type: ignore
)
data = []
for frame in container.decode(audio=0):
frame = cast(av.audio.AudioFrame, frame)
resampled_frames = resampler.resample(frame)
for resampled_frame in resampled_frames:
frame_array = resampled_frame.to_ndarray()[0]
Expand Down
18 changes: 10 additions & 8 deletions renumics/spotlight/io/audio.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,10 @@

import io
import os
from typing import IO, Dict, Tuple, Union
from typing import IO, Dict, Tuple, Union, cast

import av
import av.audio
import numpy as np
import requests
import validators
Expand Down Expand Up @@ -82,6 +83,7 @@ def read_audio(file: FileType) -> Tuple[np.ndarray, int]:

data = []
for frame in container.decode(audio=0):
frame = cast(av.audio.AudioFrame, frame)
frame_array = frame.to_ndarray()
if len(frame_array) == 1:
frame_array = frame_array.reshape((-1, num_channels))
Expand Down Expand Up @@ -123,13 +125,13 @@ def write_audio(
# `AudioFrame.from_ndarray` expects an C-contiguous array as input.
data = np.ascontiguousarray(data)
num_channels = len(data)
frame = av.audio.AudioFrame.from_ndarray(data, data_format, num_channels)
frame = av.audio.AudioFrame.from_ndarray(data, data_format, num_channels) # type: ignore
frame.rate = sampling_rate
with av.open(file, "w", format_) as container:
stream = container.add_stream(codec, sampling_rate)
stream.channels = num_channels
container.mux(stream.encode(frame))
container.mux(stream.encode(None))
stream.channels = num_channels # type: ignore
container.mux(stream.encode(frame)) # type: ignore
container.mux(stream.encode(None)) # type: ignore


def transcode_audio(
Expand All @@ -155,10 +157,10 @@ def transcode_audio(

for frame in input_container.decode(input_stream):
frame.pts = None
for packet in output_stream.encode(frame):
for packet in output_stream.encode(frame): # type: ignore
output_container.mux(packet)

for packet in output_stream.encode(None):
for packet in output_stream.encode(None): # type: ignore
output_container.mux(packet)


Expand All @@ -169,7 +171,7 @@ def get_format_codec(file: FileType) -> Tuple[str, str]:
file = prepare_input_file(file)
with av.open(file, "r") as input_container:
stream = input_container.streams.audio[0]
return input_container.format.name, stream.name
return input_container.format.name, stream.name # type: ignore


def get_waveform(file: FileType) -> np.ndarray:
Expand Down
5 changes: 3 additions & 2 deletions renumics/spotlight/media/mesh.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import io
import math
import os
from typing import IO, Dict, List, Optional, Tuple, Union
from typing import IO, Dict, List, Optional, Tuple, Union, cast
from urllib.parse import urlparse

import numpy as np
Expand Down Expand Up @@ -122,7 +122,7 @@ def from_trimesh(cls, mesh: trimesh.Trimesh) -> "Mesh":
Import a `trimesh.Trimesh` mesh.
"""
return cls(
mesh.vertices, mesh.faces, mesh.vertex_attributes, mesh.face_attributes
mesh.vertices, mesh.faces, mesh.vertex_attributes, mesh.face_attributes # type: ignore
)

@classmethod
Expand Down Expand Up @@ -155,6 +155,7 @@ def from_file(cls, filepath: PathType) -> "Mesh":
raise exceptions.InvalidFile(
f"Mesh {filepath} does not exist or could not be read."
) from e
mesh = cast(trimesh.Trimesh, mesh)
return cls.from_trimesh(mesh)

@classmethod
Expand Down

0 comments on commit 948f7c9

Please sign in to comment.