README.md

metadata

title: Crop vs No Crop Segmentation
emoji: 🌾
colorFrom: green
colorTo: yellow
sdk: docker
app_port: 8501
pinned: false

🌾 Crop vs No-Crop Segmentation

A production-ready Streamlit app that performs binary semantic segmentation with a fine-tuned NVIDIA SegFormer-B2 model, deployed on Hugging Face Spaces using the Docker SDK.

Class	Meaning	Mask color
0	Background / No Crop	Dark gray
1	Crop	Bright yellow

Inference is CPU-optimized, so it runs on free Hugging Face Spaces.

---

✨ Features

• Upload jpg / jpeg / png images.
• One-click Run Segmentation with a progress spinner.
• Three-panel output: Original · Predicted Mask · Crop Overlay (semi-transparent magenta over crop regions), all at identical size.
• Download Result — exports all three panels combined horizontally as prediction_result.png.
• Model is loaded once via @st.cache_resource.
• Weights load from a local best.pt or directly from the Hugging Face Hub.
• Automatic GPU/CPU device selection; inference under torch.no_grad().
• Friendly error handling for invalid/corrupted uploads, missing weights, and inference failures.

---

🧠 How the model is built

The architecture is rebuilt from the nvidia/mit-b2 backbone config with num_labels=2, and the fine-tuned weights are loaded on top:

MODEL_NAME = "nvidia/mit-b2"
config = SegformerConfig.from_pretrained(MODEL_NAME, num_labels=2)
model = SegformerForSemanticSegmentation(config)
model.load_state_dict(torch.load("best.pt", map_location="cpu"), strict=False)

The checkpoint is expected to be a plain state dict:

torch.save(model.state_dict(), "best.pt")

---

📂 Folder structure

crop-segmentation/
├── app.py             # Streamlit UI + workflow
├── predictor.py       # Model build (mit-b2) + best.pt loading + inference
├── utils.py           # Image loading, mask/overlay/combine, PNG export
├── requirements.txt   # Python dependencies
├── Dockerfile         # HF Spaces Docker build (Python 3.10, CPU torch)
├── README.md          # This file (with HF Spaces front matter)
├── .gitignore
├── .dockerignore
├── best.pt            # Fine-tuned weights (provide your own)
└── assets/            # Screenshots / static assets

---

💻 Local installation

git clone <your-repo-url>
cd crop-segmentation

python -m venv .venv
source .venv/bin/activate          # Windows: .venv\Scripts\activate

# CPU torch (optional but recommended for parity with deployment)
pip install torch==2.11.0 torchvision==0.26.0 \
    --index-url https://download.pytorch.org/whl/cpu
pip install -r requirements.txt

# Make sure best.pt is present, or set HF_REPO_ID (see below)
streamlit run app.py

Open http://localhost:8501.

Loading weights

Two options, controlled by environment variables:

Variable	Purpose
(none)	Use a local best.pt next to the app.
HF_REPO_ID	Download best.pt from this HF model repo, e.g. your-username/crop-segformer-b2.
HF_TOKEN	Token for a private weights repo (optional).
INPUT_SIZE	Square inference size (default 512).

export HF_REPO_ID="your-username/crop-segformer-b2"
streamlit run app.py

---

🚀 Hugging Face Spaces deployment (Docker SDK)

1. Create a new Space → SDK: Docker.
2. Push these files to the Space repository.
3. Provide the weights using one of:
   • Commit best.pt to the Space (use Git LFS for large files):

     git lfs install
     git lfs track "*.pt"
     git add .gitattributes best.pt
     git commit -m "Add weights"
     git push
     

   • Load from a model repo — in the Space Settings → Variables, set HF_REPO_ID (and HF_TOKEN if private). No best.pt in the Space needed.
4. The README.md front matter pins sdk: docker and app_port: 8501, so the Space serves the Streamlit app on port 8501 automatically.

The Space will build the Docker image, install dependencies (CPU torch), and launch:

streamlit run app.py --server.port=8501 --server.address=0.0.0.0

---

🖼️ Example screenshots

Replace these placeholders with real screenshots in assets/.

Upload	Results (3 panels)
assets/screenshot_upload.png	assets/screenshot_results.png

---

⚙️ Notes & assumptions

• Preprocessing: convert to RGB → resize to INPUT_SIZE → ToTensor → ImageNet normalization (mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]). If your fine-tuning used different stats or size, set INPUT_SIZE and adjust IMAGENET_MEAN/IMAGENET_STD in predictor.py to match.
• Postprocessing: argmax over upsampled logits, resized back to the original image resolution.
• Weights are loaded with strict=False so minor key mismatches log a warning rather than crash; verify the warnings are empty for a correct checkpoint.