# In cleanlab/segmentation/summary.py or cleanlab/object_detection/summary.py

from typing import cast
from matplotlib.axes import Axes

# When dealing with plt.subplots() return value
fig, axes = plt.subplots(1, output_plots, figsize=(5 * output_plots, 5))
if output_plots == 1:
    axes_list = [cast(Axes, axes)]
else:
    axes_list = cast(List[Axes], axes) if isinstance(axes, np.ndarray) else [axes]

# Now axes_list can be safely used without type errors
for ax in axes_list:
    ax.imshow(...)  # No mypy errors