Fix IndexError in show_segmentation when fig has no axes

[debsuvra]

Bug Report: IndexError in show_segmentation function

Description

The show_segmentation function in cellpose_omni/plot.py throws an IndexError when called with a figure that has no axes.

Error Message

IndexError: list index out of range
  File "/path/to/cellpose_omni/plot.py", line 538, in show_segmentation
    ax = fig.get_axes()[0]

Root Cause

At line 538 in cellpose_omni/plot.py, the function assumes that any provided figure already has axes:

ax = fig.get_axes()[0]

However, when creating a figure with plt.figure(), no axes are automatically created, causing fig.get_axes() to return an empty list.

Reproduction

This issue occurs when calling show_segmentation with a figure created using the examples from the documentation:

import matplotlib.pyplot as plt
from cellpose_omni import plot
import numpy as np

# Create test data
img = np.random.rand(100, 100)
mask = np.random.randint(0, 5, (100, 100))
flow = np.random.rand(100, 100, 3)

# This causes the error:
fig = plt.figure(figsize=(12, 3))
plot.show_segmentation(fig, img, mask, flow, omni=True)  # IndexError!

Proposed Fix

Replace line 538 in cellpose_omni/plot.py:

# Current (problematic) code:
ax = fig.get_axes()[0]

# Proposed fix:
if fig.get_axes():
    ax = fig.get_axes()[0]
else:
    ax = fig.add_subplot(111)

Impact

This affects any user following the documentation examples who creates a figure before calling show_segmentation. It's particularly problematic in the latest developmental version.

Environment

• Omnipose version: Latest developmental
• Python version: 3.12
• Operating System: macOS

Workaround

Currently, users can work around this by:

1. Passing None as the figure parameter (let the function create its own)
2. Creating a figure with axes: fig, ax = plt.subplots()

This fix would make the function more robust and consistent with the documentation examples.