Accurate prediction of epidemics is pivotal for making well-informed decisions for the control of infectious diseases, but addressing heterogeneity in the system poses a challenge. In this study, we propose a novel modelling framework integrating the spatio-temporal heterogeneity of susceptible individuals into homogeneous models, by introducing a continuous recruitment process for the susceptibles. A neural network approximates the recruitment rate to develop a Universal Differential Equations (UDE) model. Simultaneously, we pre-set a specific form for the recruitment rate and develop a mechanistic model. Data from a COVID Omicron variant outbreak in Shanghai are used to train the UDE model using deep learning methods and to calibrate the mechanistic model using MCMC methods. Subsequently, we project the attack rate and peak of new infections for the first Omicron wave in China after the adjustment of the dynamic zero-COVID policy. Our projections indicate an attack rate and a peak of new infections of 80.06% and 3.17% of the population, respectively, compared with the homogeneous model’s projections of 99.97% and 32.78%, thus providing an 18.6% improvement in the prediction accuracy based on the actual data. Our simulations demonstrate that heterogeneity in the susceptibles decreases herd immunity for ~37.36% of the population and prolongs the outbreak period from ~30 days to ~70 days, also aligning with the real case. We consider that this study lays the groundwork for the development of a new class of models and new insights for modelling heterogeneity.