LeafSeg: Leveraging I-JEPA, ViT, U-Net, and YOLO to gain efficiency and accuracy in leaf disease segmentation.

An idea to use the world model in segmentation task to gain efficiency

Video: https://www.youtube.com/watch?v=e2R9_CGOp5o&t=185s

Plant diseases pose a serious threat to global food security by reducing crop yield, lowering agricultural quality, and making early intervention more difficult for farmers. Although artificial intelligence has shown strong potential for agricultural image analysis, accurate plant disease detection and lesion segmentation remain challenging, especially when annotated pixel-level data are limited.

Creating segmentation masks is expensive, time consuming, and often requires expert knowledge, which makes it difficult to train high performing models in realistic settings. To address this problem, this work investigates whether self-supervised pretraining and improved data acquisition strategies can make plant disease segmentation more data-efficient while also improving segmentation accuracy.

Through modern deep learning approach, this study explores how modern deep learning methods can better leverage scarce labeled data for practical plant disease analysis.

Part 1. How I created the model: My story

1.1 Motivation and research

During late November and early December 2025, I began exploring the field of computer vision more seriously. Prior to this period, most of my focus had been on large language models (LLMs). I spent time studying how transformer architectures work by replicating Andrej Karpathy’s NanoGPT / NanoChat implementation (you can check out this article I wrote: How I replicated GPT 2 from scratch: Part 4 | by Jayden Lee | Medium) and experimenting with modifications to better understand the training process. Through this work, I learned about token embeddings, attention mechanisms, and the training dynamics of transformer based models.

While this work has been very interesting, I gradually became more curious about the fields of computer vision and thought how the similar architecture can be applied into how the model can see.

On December 19, I saw something promising while I was scrolling on Linkedin. It was a post about the JEPA style model, and how that is different from common supervised learning methods. It explained how and why JEPA style learning learned the overall representation of space, instead of just pixel-values. It predicts what it means, not what it looks like. For example, say you cover some part of the image and get the model to predict what the part of the image would look like. The normal generative model would predict the exact pixel values, while the I-JEPA style model would predict the vectorized representation of the image.

1.2. Actually learning about computer vision

To build the theoretical foundations for my project, I enrolled in Andrew Ng’s Deep Learning Specialization during the winter break. The early modules covered forward propagation, loss functions and gradient descent, reinforcing my understanding of how neural networks are trained. Subsequent modules introduced convolutional neural networks (CNNs) for processing visual data, demonstrating how convolutional filters detect spatial patterns and how pooling layers reduce dimensionality.

Furthermore, I also learned about Vision Transformers. Vision transformers (ViT) divide an image into small patches that are treated like tokens in natural language processing. Multi head self attention allows the model to capture long range relationships across the entire image. CLS token allows model to maintain comprehensive visual information for response generation. In short, Vision transformer copy paste of what I know about transformer architecture for textual tokens, only that we are using image tokens and CLS tokens.

1.3. Going deep into I-JEPA

I first learned about I-JEPA in a very natural way: not through a textbook or a formal class, but from a post I came across while scrolling on LinkedIn. That post caught my attention because it described a JEPA-style model as something fundamentally different from ordinary supervised learning and from models that try to reconstruct exact pixels. What stood out to me was the idea that I-JEPA predicts what an image region means in representation space, rather than simply recreating what it looks like pixel by pixel. At that stage, I was still new to the field, so I did not fully understand the mathematics or the architecture yet, but the core idea stayed with me. I saved it because I could tell it was important, and later that moment became the spark that pushed me to seriously explore self-supervised computer vision.

As I built a stronger foundation during winter break, I returned to I-JEPA with much more purpose. I studied deep learning, convolutional neural networks, U-Net, and Vision Transformers first, so that I could understand where I-JEPA fit into the bigger picture. Then, after reading the I-JEPA paper and related tutorials, I began to understand its main mechanism: a student network uses visible context to predict the latent representations of masked target blocks produced by a teacher network, instead of reconstructing raw pixels. That shift helped me see why I-JEPA could learn more semantic and transferable visual features, especially when labeled data are limited. By January and February, my understanding had moved beyond simple curiosity; I was now studying the teacher-student setup, masking strategies, and representation learning deeply enough to shape my own research question.

1.4. Connecting the dots

At first, the ideas I was learning, transformers, U-Net, YOLO, and self-supervised learning, felt like separate pieces. I had already spent time understanding transformers through language models, so when I began studying computer vision, I was naturally drawn to Vision Transformers and to the question of how attention-based models could understand images instead of text. Around the same time, I learned how U-Net could perform pixel-level segmentation and how YOLO could detect objects efficiently in real-world settings. The turning point came when I encountered I-JEPA and realized that it approached learning differently from standard supervised or generative methods. Instead of predicting exact pixels, it learned to predict meaningful representations from context. That idea stood out to me because plant disease images often contain complex patterns, unclear lesion boundaries, and limited labeled data, which made representation learning feel especially relevant to the problem I wanted to solve.

Once I understood these ideas more deeply, the overall direction of my project became much clearer. U-Net provided a strong segmentation framework, Vision Transformers offered global context through self-attention, and I-JEPA suggested a way to pretrain a model on unlabeled plant images before fine-tuning it on a smaller labeled segmentation dataset. At the same time, my growing interest in climate change and environmental monitoring gave the project a meaningful real-world purpose: building AI tools that could support earlier and more precise plant disease detection. In other words, the project was not built from one single idea, but from connecting several lines of learning into one system. That is how I arrived at the central design of LeafSeg: an I-JEPA-pretrained ViT encoder integrated into a U-Net-style segmentation pipeline, with YOLO supporting the broader goal of real-world plant detection and deployment.

Part 2. Intuition and technicality

Now, we get into the technicality of my LeafSeg project and how it is structured.

LeafSeg is built as a two-stage vision system centered on a shared Vision Transformer backbone. In the first stage, the encoder is pretrained with I-JEPA, where the model does not reconstruct missing pixels, but instead predicts the latent representations of masked image regions from visible context. The context encoder is a standard ViT, and its job is to process the patches of the visible patch of an image that are given to model as a clue. To obtain the context in I-JEPA, we first sample a single block x from the image with a random scale in the range (0.85, 1.0) and unit aspect ratio. To ensure a non-trivial prediction task, we remove any overlapping regions from the context block. The masked context block, x, is fed through the context encoder to obtain a corresponding patch-level representation.

The target encoder uses exponential weighted average, so it is not trained explicitly with gradient descent. Rather, it is trained with the student encoder to prevent model from representation collapse. Given an input image y, we convert it into a sequence of N non overlapping patches, and feed this through the target-encoder to obtain a corresponding patch level representation s of the image. We obtain the loss by sampling M, possibly overlapping, blocks from the target representations s. Typically, we set M equal to 4, and sample the blocks with a random aspect ratio in the range (0.75,1.5) and random scale in the range (0.15,0.2).

The predictor is also a ViT and it takes two things as an input: the output representation for the context encoder and the positional mask tokens. The positional mask tokens tell the predictor where in the image the target block is located. The mask token is parameterized by a shared learnable vector with an added positional embedding. Since we wish to make predictions for M target blocks, we apply our predictor M times. As a result, the predictor outputs a patch level predicted representation for that specific target block.

The loss is simply the average L2 distance between the predicted patch-level embedding and target representation, given by this formula:

Picture from 2301.08243

For LeafSeg, most of the hyperparameter settings match with the original I-JEPA implementation. It has four prediction masks, (0.85, 1.00) encoder mask scale, and (0.15, 0.20) predictor mask scale.

Then, pretrained ViT encoder (the context encoder) is then reused in the second stage, where it feeds a U-Net-style decoder for pixel-level lesion segmentation. The U-Net style decoder has down blocks, bottleneck, upsampling path, skip connections, and final 1x1 output layer. Technically, this means LeafSeg is not simply a standard U-Net with a transformer inserted into it; it is better understood as an I-JEPA-pretrained ViT semantic backbone followed by a spatial decoder that reconstructs dense lesion masks. If you want to learn more about U-Net, follow this article: The State-of-the-art model in Input Segmentation and Object Detection | by Jayden Lee | Feb, 2026 | Medium

U-Net style decoder is needed because pure ViT token outputs are semantically strong but spatially coarse. Segmentation needs both global context and local boundary recovery. The U-Net decoder is there to restore spatial detail and shape continuity. In this design, the ViT answers “what is present and where roughly,” while the U-Net decoder answers “what are the precise mask boundaries.”

Furthermore, this architecture uses GroupNorm because this repo is intended for more small-size dataset training. Small batch segmentation often makes batchnorm unstable or noisy, so groupnorm is a better fit for this repo’s intended training environment.

In summary, the intuition behind this design is that plant disease segmentation requires two different kinds of visual intelligence at once. First, the model must understand broad structure: what part of the image is leaf, what patterns look pathological, and how lesion regions relate to surrounding healthy tissue. That is where the ViT and I-JEPA pairing matters, because the transformer processes the image as patches and learns long-range relationships, while I-JEPA encourages it to predict meaning in representation space rather than memorize texture or raw pixel appearance. Second, the model must recover precise lesion boundaries, which is why the U-Net-style decoder is still necessary. The transformer provides semantically rich but spatially coarse features, and the decoder refines those features into detailed masks. In simple terms, the ViT answers “what is happening in this image and roughly where,” while the decoder answers “what are the exact diseased pixels.” This is also why the architecture is especially suited to limited-label settings: the encoder can first learn useful visual structure from unlabeled plant images, then transfer that knowledge to segmentation, which matches your report’s conclusion that I-JEPA improved generalization, label efficiency, and convergence speed compared with the supervised alternatives.

For real-world deployment, LeafSeg extends beyond segmentation alone. I use YOLO as the front-end detector that first localizes plants or leaves in messy real environments, after which the segmentation model can focus on diseased regions more precisely. That makes the full pipeline practical: YOLO handles fast plant detection, while the I-JEPA–ViT–U-Net handles fine-grained disease localization. Conceptually, this creates a “detect first, segment second” system that is much closer to how an agricultural monitoring tool would work outside the lab.

Part 3. The hypothesis, method, and result

Hypothesis: A Vision Transformer encoder pre-trained using the I-JEPA self-supervised method and integrated into a U-Net decoder will achieve significantly higher plant disease lesion segmentation and classification performance than baseline models when trained on the same limited labeled dataset.

Now that the hypothesis and the method for the model is identified, I introduce the method to create the experiment:

1. First, create the proposed I-JEPA ViT U-Net model with pytorch
2. Create the Baseline U-Net model
3. Compare on the Plant Disease Dataset (PlantSeg and Leaf Disease Segmentation dataset)

These methods will help to design the architecture and test its compatibility.

The engineering objective of this project is to develop a robust deep-learning pipeline for automated plant disease diagnosis that can operate in real-world agricultural environments.

The system is designed to:

• Detect plants or leaves within an image using an object detection model (YOLO)
• Segment diseased regions at the pixel level using a ViT-U-Net architecture

The pipeline aims to achieve high segmentation and classification performance while trying out different architecture to see which one is even slightly better.

For the dataset, I used PlantSeg Dataset and Leaf Disease Segmentation dataset. They are two publicly available plant disease datasets, and they were used for training and evaluation.

PlantSeg dataset

The PlantSeg dataset contains images of plant leaves paired with pixel-level segmentation masks that mark diseased regions. These annotations allow deep learning models to learn precise localization of plant disease lesions, rather than simply classifying the entire leaf as healthy or diseased. The dataset is widely used to train and evaluate segmentation models such as U-Net or transformer-based architectures for automated plant disease detection.

Leaf Disease Segmentation dataset

The Leaf Disease Segmentation dataset contains 7,056 images with pixel-level annotations of diseased leaf regions, making it suitable for training and evaluating semantic segmentation models. Each image is paired with a ground-truth mask that highlights infected areas such as lesions, spots, or blight, allowing models to learn precise localization of plant disease symptoms. This dataset is commonly used to benchmark segmentation architectures like U-Net and transformer-based models for automated plant disease detection in agricultural monitoring systems.

Results from the held-out test sets showed that the proposed I-JEPA–ViT–U-Net outperformed the comparison models. With just a 20 minutes of training, it achieved the validation Dice score of 0.62 and mIou score of 0.67. Training further shows that I-JEPA-ViT-U-Net achieves the best overall segmentation performance with a Dice coefficient of 0.70 and an mIoU of 0.74. Relative to both the baseline U-Net and the supervised ViT–U-Net, the self-supervised model demonstrated improved lesion localization, as reflected by higher Dice and IoU. The label-fraction experiments further showed that the I-JEPA-pretrained encoder was especially beneficial when labeled data were scarce, where it maintained a small but consistent advantage over the supervised alternative. Visual comparisons supported these quantitative findings, showing finer boundary recovery and fewer false detections in healthy tissue, while the training curves suggested faster convergence during fine-tuning. For the real-world detection component, YOLOv11 achieved a loss of 0.9265 and it generalized reliably in plant detection, strengthening the feasibility of the full LeafSeg system for practical agricultural use.

Picture of I-JEPA and scratch comparison

Picture of Dice score comparsion between I-JEPA and Scratch baseline (no I-JEPA pretraining)

Picture of comparison between scratch baseline and ijepa baseline

Part 4. Discussions

Although the proposed model achieves high accuracy, several limitations remain:

• Computational cost: ViT encoders require significant computational resources. Pre‑training I‑JEPA on large datasets demands multiple GPUs and substantial training time. Also, there are some overfitting issues that I have found. This may be prohibitive for researchers with limited resources.
• Dataset bias: The plant disease datasets used in this study may not encompass all possible disease patterns, species or environmental conditions. Domain shifts (e.g., different camera types, backgrounds, field conditions) could affect performance. Evaluating on additional datasets and in real field settings is necessary.
• Annotation quality: Segmentation masks in publicly available datasets may have noise and inconsistencies. Better annotation protocols and quality control would improve training and evaluation.
• Edge deployment: ViT models are large, raising challenges for deployment on resource‑constrained IoT devices. Model compression, pruning and quantization should be explored to enable real‑time inference.

Summary: This research investigated the impact of self‑supervised I‑JEPA pre‑training on ViT encoders for leaf disease segmentation and classification. We integrated an I‑JEPA‑pretrained ViT encoder into a U‑Net decoder and compared it with baseline U‑Net and supervised ViT–U‑Net models. Results on the PlantSeg and Leaf Disease Segmentation datasets show that the proposed model achieves a higher score compared to classical U-Net but challenges remain. Statistical analysis confirms that these improvements are significant. The findings suggest that I‑JEPA pre‑training enhances feature representations, accelerates convergence and improves generalization when labeled data are limited. Thus, self‑supervised ViT encoders represent a promising direction for precision agriculture applications.

Furthermore, the pipeline for YOLOv11 and I-JEPA ViT U-Net model are well established and can be deployed to real life settings.

Author’s bio: Author’s Bio: My name is Jayden Lee, and thank you for reading this article. I’m deeply curious about how artificial intelligence learns, perceives patterns, and builds understanding through modern deep learning systems. I’m particularly fascinated by the mechanics behind training neural networks and the ideas that connect intuition, mathematics, and computation. I believe many real-world problems can be addressed through vision–language models and deep learning, and I hope to be part of the generation that pushes this field forward. If you’d like to connect or discuss AI, feel free to reach out to me on LinkedIn: Jayden Lee | LinkedIn

Sources:

Assran, M., Duval, Q., Misra, I., Bojanowski, P., Vincent, P., Rabbat, M., LeCun, Y., & Ballas, N. (2023). Self-supervised learning from images with a joint-embedding predictive architecture. arXiv. https://arxiv.org/abs/2301.08243

Rath, S. R. (2025, August 11). JEPA series part 1: Introduction to I-JEPA. DebuggerCafe. https://debuggercafe.com/jepa-series-part-1-introduction-to-i-jepa/

Meta AI. (n.d.). IJEPA: Official codebase for I-JEPA [Computer software]. GitHub. https://github.com/facebookresearch/ijepa