Computer Vision for Stressed Crop Detection

This repository contains tools and deep learning models for stressed crop detection through image segmentation.

Directory Structure

• src/: Contains all the source code for data preparation, model training, and inference.
  • main.py: Entry point for training and evaluation.
  • inference.py: Script to run inference on new images.
  • cv_masking.py: OpenCV script for generating yellow color masks.
  • network.py: PyTorch models (U-Net, R2U-Net, AttU-Net, R2AttU-Net).
  • dataset.py & data_loader.py: Utilities for preparing datasets.
  • solver.py: Core logic for training and evaluation loop.
• weights/: Contains pre-trained PyTorch model weights (.pt or .pkl).

Setup

Ensure you have PyTorch, torchvision, and OpenCV installed:

pip install torch torchvision opencv-python Pillow numpy

Usage

1. Color Masking (Data Prep)

Generate a yellow-color mask (e.g., for stressed regions) using OpenCV:

python src/cv_masking.py --input path/to/images --output ./masks

2. Training

To train the segmentation models, arrange your dataset into a directory and run:

python src/main.py --mode train --model_type U_Net --train_path ./dataset/train/ --valid_path ./dataset/valid/

(Available models: U_Net, R2U_Net, AttU_Net, R2AttU_Net)

3. Inference

Run inference on new images using a trained model weight file:

python src/inference.py --input path/to/image_or_folder --output ./results --model_type U_Net --weights weights/model_checkpoint1.pt

Credits

Based on image segmentation architectures such as U-Net and its attention/recurrent variants.