| | --- |
| | license: apache-2.0 |
| | tags: |
| | - vision |
| | --- |
| | |
| |
|
| | # VisualSplit |
| |
|
| | **VisualSplit** is a ViT-based model that explicitly factorises an image into **classical visual descriptors**—such as **edges**, **color segmentation**, and **grayscale histogram**—and learns to reconstruct the image conditioned on those descriptors. This design yields **interpretable representations** where geometry (edges), albedo/appearance (segmented colors), and global tone (histogram) can be reasoned about or varied independently. |
| |
|
| | > **Training data**: ImageNet-1K. |
| | --- |
| |
|
| | ## Model Description |
| |
|
| | - **Inputs** (at inference): |
| | - An RGB image (for convenience) which is converted to descriptors using the provided `FeatureExtractor` (edges, color segmentation, grayscale histogram). |
| | - **Outputs**: |
| | - A reconstructed RGB image tensor (same spatial size as the model’s training resolution; default `224×224` unless you trained otherwise). |
| |
|
| | --- |
| |
|
| | ## Getting Started (Inference) |
| |
|
| | Below are two ways to run inference with the uploaded `model.safetensors`. |
| |
|
| | ### 1) Minimal PyTorch + safetensors (load state dict) |
| |
|
| | ```python |
| | import torch |
| | from huggingface_hub import hf_hub_download |
| | from safetensors.torch import load_file |
| | |
| | # 1) Import your model & config from the VisualSplit repo |
| | from visualsplit.models.CrossViT import CrossViTForPreTraining, CrossViTConfig |
| | from visualsplit.utils import FeatureExtractor |
| | |
| | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| | |
| | # 2) Build a config matching your training (edit if you changed widths/depths) |
| | config = CrossViTConfig( |
| | image_size=224, # change if your training size differs |
| | patch_size=16, |
| | # ... any other config fields your repo exposes |
| | ) |
| | |
| | model = CrossViTForPreTraining(config).to(device) |
| | model.eval() |
| | |
| | # 3) Download and load state dict from this model repo |
| | # Replace REPO_ID with your Hugging Face model id, e.g. "HenryQUQ/visualsplit") |
| | ckpt_path = hf_hub_download(repo_id="REPO_ID", filename="model.safetensors") |
| | state_dict = load_file(ckpt_path) |
| | missing, unexpected = model.load_state_dict(state_dict, strict=False) |
| | print("Missing keys:", missing) |
| | print("Unexpected keys:", unexpected) |
| | |
| | # 4) Prepare an input image and extract descriptors |
| | from PIL import Image |
| | from torchvision import transforms |
| | |
| | image = Image.open("input.jpg").convert("RGB") |
| | transform = transforms.Compose([ |
| | transforms.Resize((config.image_size, config.image_size)), |
| | transforms.ToTensor(), |
| | ]) |
| | pixel_values = transform(image).unsqueeze(0).to(device) # (1, 3, H, W) |
| | |
| | # FeatureExtractor provided by the repo should return the required tensors |
| | extractor = FeatureExtractor().to(device) |
| | with torch.no_grad(): |
| | edge, gray_hist, segmented_rgb, _ = extractor(pixel_values) |
| | |
| | # 5) Run inference (reconstruction) |
| | with torch.no_grad(): |
| | outputs = model( |
| | source_edge=edge, |
| | source_gray_level_histogram=gray_hist, |
| | source_segmented_rgb=segmented_rgb, |
| | ) |
| | # Your repo’s forward returns may differ; adjust the key accordingly: |
| | reconstructed = outputs["logits_reshape"] # (1, 3, H, W) |
| | |
| | # 6) Convert to PIL for visualisation |
| | to_pil = transforms.ToPILImage() |
| | recon_img = to_pil(reconstructed.squeeze(0).cpu().clamp(0, 1)) |
| | recon_img.save("reconstructed.png") |
| | print("Saved to reconstructed.png") |
| | ``` |
| |
|
| | ### 2) Reproducing the notebook flow (`notebook/validation.ipynb`) |
| |
|
| | The repository provides a validation notebook that: |
| | 1. Loads the trained model, |
| | 2. Uses `FeatureExtractor` to compute **edges**, **color-segmented RGB**, and **grayscale histograms**, |
| | 3. Runs the model to obtain a reconstructed image, |
| | 4. Saves/visualises the result. |
| |
|
| | --- |
| |
|
| | ## Installation & Requirements |
| |
|
| | ```bash |
| | # clone the VisualSplit code |
| | git clone https://github.com/HenryQUQ/VisualSplit.git |
| | cd VisualSplit |
| | # pip install -e . |
| | ``` |
| |
|
| | --- |
| |
|
| | ## Training Data |
| |
|
| | - **Dataset**: **ImageNet-1K**. |
| | - |
| | > This repository only hosts the **trained checkpoint for inference**. Follow the GitHub repo for the full training pipeline and data preparation scripts. |
| |
|
| | --- |
| |
|
| | ## Model Sources |
| |
|
| | - **Code**: https://github.com/HenryQUQ/VisualSplit |
| | - **Weights (this page)**: this Hugging Face model repo |
| |
|
| | --- |
| |
|
| | ## Citation |
| |
|
| | If you use this model or ideas, please cite: |
| |
|
| | ```bibtex |
| | @inproceedings{Qu2025VisualSplit, |
| | title = {Exploring Image Representation with Decoupled Classical Visual Descriptors}, |
| | author = {Qu, Chenyuan and Chen, Hao and Jiao, Jianbo}, |
| | booktitle = {British Machine Vision Conference (BMVC)}, |
| | year = {2025} |
| | } |
| | ``` |
| |
|
| | --- |
