Makes summary writer support individual images as well.
Makes summary writer support individual images as well.

Note that `metric_writers.interface.MetricWriter.write_images()` says both shapes `[H, W, C]` and `[N, H, W, C]` are supported.

PiperOrigin-RevId: 394178318

google/CommonLoopUtils,google/CommonLoopUtils