README.md
3,928 bytes
| 1 | # CountrySense |
|---|---|
| 2 | |
| 3 | Guess the country from a single street-level photo. A fine-tuned CLIP vision transformer classifies images into ~100+ countries. Country classification only: no coordinates, no city guessing. |
| 4 | |
| 5 | Estonian documentation: [README.et.md](README.et.md) (full write-up) and [docs/mudel.html](docs/mudel.html) (layer-by-layer explanation of the architecture, data selection and cleaning). |
| 6 | |
| 7 | ## Why country-level, not coordinates |
| 8 | |
| 9 | Exact geolocation (predicting latitude and longitude) is research-grade work: it needs millions of images, large models and compute budgets this project does not have. Country classification over ~100 countries is a tractable supervised problem: fine-tune a pretrained vision model on a free Colab or Kaggle GPU in an afternoon and get accuracy that feels surprising in a demo. Region-level prediction within a country is the honest next step on the roadmap. Coordinates are out of scope. |
| 10 | |
| 11 | ## Model |
| 12 | |
| 13 | - Backbone: ViT-B/16 pretrained by OpenAI CLIP (`vit_base_patch16_clip_224.openai` via timm), 86M parameters. |
| 14 | - Head: `Dropout(0.2)` + a single `Linear(768 -> num_countries)` layer. |
| 15 | - Why CLIP: it was pretrained on 400M image-text pairs from the web, so its features already encode things like signage, vegetation, architecture and road furniture. A linear probe on CLIP features is a strong geolocation baseline; fine-tuning the last blocks improves on it. |
| 16 | |
| 17 | Training recipe (fits on a free T4, roughly 1 to 2 hours end to end): |
| 18 | |
| 19 | 1. Phase 1, linear probe: backbone frozen, train only the head, 3 epochs, lr 1e-3. This protects pretrained weights from gradients of a randomly initialized head. |
| 20 | 2. Phase 2, partial fine-tune: unfreeze the last 4 transformer blocks and the final norm, lr 2e-5 (backbone) / 1e-4 (head), cosine schedule, 8 epochs. |
| 21 | |
| 22 | Cross-entropy with label smoothing 0.1, AdamW, AMP, gradient clipping, `WeightedRandomSampler` for class imbalance. No horizontal flip augmentation: driving side (left vs right traffic) is a real geographic cue that flipping would corrupt. |
| 23 | |
| 24 | ## Data |
| 25 | |
| 26 | Default dataset: [GeoLocation, Geoguessr Images 50K](https://www.kaggle.com/datasets/ubitquitin/geolocation-geoguessr-images-50k) on Kaggle (~50k Street View images in per-country folders). Any dataset laid out as `root/<country>/<image>` works. |
| 27 | |
| 28 | Cleaning (`countrysense/clean.py`) drops corrupt files, images smaller than 128px, too dark / too bright / blurry images, perceptual-hash duplicates, and countries with fewer than 100 images. It writes a manifest CSV plus a per-image report of what was dropped and why. |
| 29 | |
| 30 | ## Quickstart |
| 31 | |
| 32 | Colab (recommended): open `notebooks/train_colab.ipynb`, select a T4 GPU runtime and run all cells. |
| 33 | |
| 34 | Local: |
| 35 | |
| 36 | ```bash |
| 37 | pip install -r requirements.txt |
| 38 | python -m countrysense.clean --raw-dir data/raw |
| 39 | python -m countrysense.train --config configs/default.yaml |
| 40 | python -m countrysense.evaluate --checkpoint runs/clip_vit_b16/best.pt |
| 41 | python -m countrysense.predict --image path/to/photo.jpg |
| 42 | ``` |
| 43 | |
| 44 | Optional Gradio demo: `pip install gradio` and `python app.py`. |
| 45 | |
| 46 | ## Evaluation |
| 47 | |
| 48 | `evaluate.py` reports top-1, top-5 and macro F1, writes per-country accuracy and a confusion matrix. Top-5 matters for geography: confusing Estonia with Latvia is a much smaller error than confusing it with Brazil, and neighbor confusion dominates the mistakes. |
| 49 | |
| 50 | ## Layout |
| 51 | |
| 52 | ``` |
| 53 | countrysense/ package: data, clean, model, train, evaluate, predict |
| 54 | configs/default.yaml training configuration |
| 55 | notebooks/ Colab notebook, end to end |
| 56 | docs/mudel.html architecture explainer (Estonian) |
| 57 | app.py Gradio demo |
| 58 | ``` |
| 59 | |
| 60 | ## Roadmap |
| 61 | |
| 62 | - [ ] Region-level prediction within a country (hierarchical head: country, then region) |
| 63 | - [ ] Larger training set (OpenStreetView-5M subset) |
| 64 | - [ ] Calibration and abstention ("not confident enough to guess") |
| 65 | - Not planned: exact coordinates. That is research territory (see PIGEON, GeoCLIP) and needs compute this project intentionally avoids. |
| 66 | |