This work presents a framework that jointly optimizes a diffractive optical system and a segmentation network for spatio-spectral classification. The diffractive optical element (DOE) is learned jointly with the network via a differentiable light-propagation model, thereby shaping the captured measurements to emphasize class-separating features under a reduced number of sensor channels. The main contribution is a spectral separability regularizer based on class-wise mean spectral signatures that increases the distance between the average class signatures in the captured measurements. For each training image, one representative signature per class is obtained by averaging the measured spectra over the pixels of that class, and the regularizer encourages these class-wise mean signatures to be farther apart. The goal is separability at acquisition, since the jointly optimized optics shape the captured multiband data so that small spectral differences between classes become easier to observe by the spatial classification network across the available sensor channels, embedding class-distinctive cues directly in the measurements. This reduces the need for many spectral bands and points to simpler and lower-cost sensors without sacrificing classification quality. The study uses unmanned aerial vehicle (UAV) hyperspectral imaging (HSI) data and evaluates three spectral-band settings (10, 100, and 200 channels), each tested with a traditional refractive lens and with the proposed jointly optimized DOE. Across all settings the optimized optics improve segmentation accuracy, with the largest gains in the 10-channel case. The learned optics in the 10-channel setting achieve accuracy close to the 200-channel baseline while relying on far fewer sensor channels. Ablation experiments examine phase initialization and measure the effect of the spectral separability regularizer. The results support task-driven optical design as a compact and practical route to high-quality spatio-spectral classification when only a limited number of sensor channels are available.
Code release notice: The code required to reproduce our results will be made publicly available upon acceptance of the paper.