Update README.md

README.md

@@ -5,22 +5,23 @@ tags:
 - unet
 - pytorch
 license: apache-2.0
+datasets:
+- johnowhitaker/imagenette2-320
 ---
-
 # 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
+A 3,968,892-parameter U-Net that colorizes grayscale photos. Classification-style  
+(Zhang et al., *Colorful Image Colorization*, [arXiv:1603.08511](https://arxiv.org/abs/1603.08511)): 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
+- **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
 
@@ -31,16 +32,12 @@ 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)
-```
+rgb_out = np.clip(lab2rgb(lab_out), 0, 1)