The WKBJ solution for the one-way wave equations in media with smoothly varying velocity variation with depth, c(z), is reformulated from the principle of energy flux conservation for acoustic media. The formulation is then extended to general heterogeneous media with local angle domain methods by introducing the concepts of Transparent Boundary Condition (TBC) and Transparent Propagator (TP). The influence of the WKBJ correction on image amplitudes in seismic imaging, such as depth migration in exploration seismology, is investigated in both smoothly varying c(z) and general heterogeneous media. We also compare the effect of the propagator amplitude compensation with the effect of the acquisition aperture correction on the image amplitude. Numerical results in a smoothly varying c(z) medium demonstrate that the WKBJ correction significantly improves the one-way wave propagator amplitudes, which, after compensation, agree very well with those from the full wave equation method. Images for a point scatterer in a smoothly varying c(z) medium show that the WKBJ correction has some improvement on the image amplitude, though it is not very significant. The results in a general heterogeneous medium (2D SEG/EAGE salt model) show similar phenomena. When the acquisition aperture correction is applied, the image improves significantly in both the smoothly varying c(z) medium and the 2D SEG/EAGE salt model. The comparisons indicate that although the WKBJ compensation for propagator amplitude may be important for forward modeling (especially for wide-angle waves), its effect on the image amplitude in seismic imaging is much less noticeable compared with the acquisition aperture correction for migration with limited acquisition aperture in general heterogeneous media.