README.md

---
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:

```python
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:

```python
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

```bash
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`). |

```bash
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):
     ```bash
     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:

```bash
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.