Let $F/F^+$ be a CM field and let $\widetilde{v}$ be a finite unramified place of $F$ above the prime $p$. Let $\overline{r}: \mathrm{Gal}(\overline{\mathbb{Q}}/F)\rightarrow \mathrm{GL}_n(\overline{\mathbb{F}}_p)$ be a continuous representation which we assume to be modular for a unitary group over $F^+$ which is compact at all real places. We prove, under Taylor--Wiles hypotheses, that the smooth $\mathrm{GL}_n(F_{\widetilde{v}})$-action on the corresponding Hecke isotypical part of the mod-$p$ cohomology with infinite level above $\widetilde{v}|_{F^+}$ determines $\overline{r}|_{\mathrm{Gal}(\overline{\mathbb{Q}}_p/F_{\widetilde{v}})}$, when this latter restriction is Fontaine--Laffaille and has a suitably generic semisimplification.