Smile Design with AI: from Tooth Segmentation to Whitening and Alignment

Overview

This project is a smile-design preview system for cosmetic dentistry. A patient gives an ordinary photo of a smile, and in a few seconds the system shows how the teeth would look after whitening, shading and straightening, while the rest of the face stays the same. The aim is that the dentist and the patient reach a shared picture of the result before any treatment. I built this project in 2025, and a demo of it runs at medical.alipour.me.

A large part of the work was research and experiments: for every step I read several papers, built several methods and compared them. On this page I explain all the methods I tried and how I built the system. The pictures on this page come from my own experiments and test data.

Results

In the pictures at the top, the original photo is on the left and the output of the system is on the right. The system changes only the teeth, and the rest of the face, skin, hair and background stay untouched.

What the problem was

The teeth must be found in a photo taken with a phone, cut out exactly (the borders of neighbouring teeth, the gum and the lip), their colour and texture changed naturally, and, when needed, crooked teeth straightened, without making the face look "artificial". Each of these is a separate problem and I tried several methods for each.

Step 1: finding the mouth

• Face parsing: with the BiSeNet model, which splits every pixel of the face into 19 parts (lips, skin, teeth, hair and so on), I find the mouth region and cut it out with a percentage margin, so the work runs on a small, detailed image. After processing, the result is put back with a soft blend.

• Before that, I also tried older methods such as the 68 landmarks of dlib, MediaPipe Face Mesh and classic mouth detection. For tilted photos I calculate the angle of the mouth from the mask and rotate the mask.

Step 2: cutting out the teeth (segmentation)

This step took the most time, because everything depends on an exact mask of the teeth. Pictures 7 and 8 show the output of this step: every tooth with its own border, even crooked and overlapping ones.

• Building the data. I prepared more than 700 photos of teeth with polygon labels (COCO format) and wrote tools to convert labels to masks, remove mismatched pairs, make the sizes equal and augment the data. With augmentation the training set reached about 5,000 images, with separate validation and test sets.

• U²-Net and U²-Net-P. I trained a U-Net style segmentation network from scratch, both the full and the light version, with precision, recall and IoU metrics and loss curves. I tried it with one class (teeth) and with several classes (individual teeth).

• YOLO11 in segmentation mode. I trained and fine-tuned the n, m and x versions on the same data, and also exported the output to ONNX. YOLO is fast and finds every tooth separately as an object.

• SAM (Segment Anything) and SAM-HQ. The box of every tooth from YOLO is given to SAM to get the exact border of the tooth. I also tried its light mobile version for speed. The combination "YOLO finds, SAM cuts exactly" gives more exact borders than either one alone.

• Segmentation papers. I read and tested SegmentAnyTooth and the individual tooth segmentation with pseudo edge-region, and took ideas from them for counting and numbering the teeth.

Step 3: colour and texture

• Colour correction (picture 13). The teeth are analysed in the HLS and HSV colour spaces: first the yellow is reduced, then the gum-like red that falls on the edge of the tooth, with adjustable ranges of hue and saturation. Because it works only inside the mask, the lips and the gums are not coloured.

• Enamel texture. A real tooth texture (with translucent edges and a soft shadow) is laid over the tooth with a faded mask and several blend modes, so that the tooth does not look flat and "plastic". Natural detail and brightness are then restored.

• Central incisors. Because the eye is sensitive to the symmetry of the two front teeth, I measure their proportion (for example the ratio of the left and right tooth) and correct the shape.

• For every step I made presets that the dentist can change, and the demo has a "server default" and a "my own settings" version.

Step 4: straightening the teeth (alignment)

Predicting how the teeth look after orthodontics in a 2D photo. I read the papers iOrthoPredictor (SIGGRAPH Asia), the 3D-structure-guided network and OrthoAligner, and implemented their approach: first the contour of the crooked teeth is turned into an aligned contour (Contour2Contour), then, from the aligned contour, the tooth image is built with the face colour and the lighting (Contour2Tooth), using a generative model of the diffusion type. Pictures 9 to 12 are the "before" and "after" of this step.

Step 5: generative methods (diffusion and inpainting)

Instead of adjusting colours, we mask the tooth area and a generative model rebuilds it. The experiments:

• Stable Diffusion XL with its inpainting version, with LoRAs, SDXL Turbo for speed, and models tuned for faces (Juggernaut, DreamShaper).

• Fooocus with Juggernaut XI, for which I built a separate Flask server that the demo calls.

• Flux (FLUX.2) with inpainting graphs in ComfyUI.

• HD-Painter (prompt-faithful, high-quality inpainting) and a fast Stable Diffusion 3.

• pix2pix. I prepared before-and-after pairs (side by side, 512 pixels) to train a conditional network.

The general difference: generative methods give the most natural texture but may change the tooth shape or the rest of the face, and are slower; the classic method (an exact mask and colour correction) is fast and controllable and keeps the face intact. For this reason I built the main path on the exact mask of step 2 and use the colour or generative method only inside the mask.

Running and delivery

• A Flask service with Docker on a CUDA base, ready to run on RunPod (serverless mode with an /invocations endpoint) and a remote GPU.

• A tunnel server and a web app to reach the user: a Next.js demo with React and Tailwind CSS, with samples, photo upload and settings.

• I also built a tool that aligns before and after photos with MediaPipe for a fair comparison.

The diagram below shows the overall path.

Tools

• Python, PyTorch, OpenCV, MediaPipe, dlib, Flask

• YOLO11 (Ultralytics), U²-Net, SAM and SAM-HQ, BiSeNet

• Stable Diffusion XL, SDXL Turbo, Stable Diffusion 3, Flux, LoRA, Fooocus, ComfyUI, HD-Painter

• Docker, CUDA, RunPod, Next.js, TypeScript

Scientific sources I read

• SegmentAnyTooth: An open-source deep learning framework for tooth enumeration and segmentation in intraoral photos. Journal of Dental Sciences, 2025.

• Individual tooth segmentation in human teeth images using pseudo edge-region obtaining. Signal Processing: Image Communication, 2024.

• iOrthoPredictor: Model-guided Deep Prediction of Teeth Alignment. ACM SIGGRAPH Asia, 2020.

• 3D Structure-guided Network for Tooth Alignment in 2D Photograph. Y. Dou et al.

• OrthoAligner: Image-Based Teeth Alignment Prediction via Latent Style Manipulation. IEEE TVCG, 2023.

• HD-Painter: High-Resolution and Prompt-Faithful Text-Guided Image Inpainting with Diffusion Models. arXiv:2312.14091.

Result

A complete system from an ordinary photo to a smile preview, built from several separate steps (finding the mouth, cutting out the teeth, colour and texture, straightening the teeth, a generative method), each of which was tried and compared with several methods.