Upload README.md with huggingface_hub

README.md

@@ -1,16 +1,46 @@
 ---
-license: apache-2.0
-pipeline_tag: image-to-image
 tags:
-- classification
-- colorization
 - image-to-image
-- model_hub_mixin
-- pytorch_model_hub_mixin
+- colorization
 - unet
+- pytorch
+license: apache-2.0
 ---
 
-This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
-- Code: [More Information Needed]
-- Paper: [More Information Needed]
-- Docs: [More Information Needed]
+# Mini U-Net Colorizer
+
+A 3,968,892-parameter U-Net that colorizes grayscale photos. Classification-style
+(Zhang et al., "Colorful Image Colorization"): it predicts a distribution over
+236 quantized CIE Lab `a`/`b` bins per pixel rather than regressing a
+single ab value directly, with color-bin loss weights derived from the real
+training-data distribution (rare/saturated colors weighted higher) so it
+doesn't just hedge toward desaturated averages. Decode with an annealed mean.
+
+- **Status:** Final (20 epochs complete)
+- **Input:** L channel normalized as `L/50 - 1` -> `[-1, 1]`, shape `(1, 256, 256)`
+- **Output:** logits over 236 ab bins, shape `(236, 256, 256)`
+- **Trained on:** `johnowhitaker/imagenette2-320` (None), warm-started from `User-2468/mini-unet-colorizer`
+- **Loss so far (weighted soft cross-entropy):** train 2.4069, val 2.7111
+
+## Usage
+
+```python
+import numpy as np, torch
+from skimage.color import rgb2lab, lab2rgb
+from PIL import Image
+# paste the SmallUNetColorizer class definition from the training script, then:
+model = SmallUNetColorizer.from_pretrained("User-2468/mini-unet-colorizer")
+model.eval()
+
+img = Image.open("photo.jpg").convert("RGB").resize((256, 256))
+lab = rgb2lab(np.asarray(img).astype("float32") / 255.0)
+L = torch.from_numpy(lab[:, :, 0:1] / 50.0 - 1.0).permute(2, 0, 1)[None]
+
+with torch.no_grad():
+    logits = model(L)
+    ab = model.decode(logits, temperature=0.38)[0].permute(1, 2, 0).numpy()
+
+L_out = (L[0, 0].numpy() + 1) * 50.0
+lab_out = np.concatenate([L_out[:, :, None], ab], axis=-1)
+rgb_out = np.clip(lab2rgb(lab_out), 0, 1)
+```