Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Correcting image white balance with Python PIL and Numpy
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": false
},
"outputs": [],
"source": [
"from PIL import Image\n",
"import numpy as np\n",
"from urllib.request import urlopen\n",
"\n",
"im = Image.open(urlopen('https://codeandlife.com/wp-content/uploads/2019/07/pistore-8-768x654.jpg'))\n",
"#im = Image.open('test.jpg') # For local files\n",
"im"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data = np.asarray(im)\n",
"data.shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"refpos = (550,360) # top left corner\n",
"refsize = (32,32) # reference sample size\n",
"sub = data[refpos[0]:refpos[0]+refsize[0],refpos[1]:refpos[1]+refsize[1]]\n",
"Image.fromarray(sub)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"c = list(np.mean(sub[:,:,i]) for i in range(3))\n",
"c"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"wb = data.astype(float)\n",
"\n",
"# RGB scaling just decreases oversaturated components to get average grey\n",
"# See https://en.wikipedia.org/wiki/Color_balance\n",
"\n",
"for i in range(3): wb[:,:,i] /= c[i]/float(min(c))\n",
" \n",
"Image.fromarray(wb.astype(np.uint8))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Conversion functions courtesy of https://stackoverflow.com/a/34913974/2721685\n",
"def rgb2ycbcr(im):\n",
" xform = np.array([[.299, .587, .114], [-.1687, -.3313, .5], [.5, -.4187, -.0813]])\n",
" ycbcr = im.dot(xform.T)\n",
" ycbcr[:,:,[1,2]] += 128\n",
" return ycbcr #np.uint8(ycbcr)\n",
"\n",
"def ycbcr2rgb(im):\n",
" xform = np.array([[1, 0, 1.402], [1, -0.34414, -.71414], [1, 1.772, 0]])\n",
" rgb = im.astype(np.float)\n",
" rgb[:,:,[1,2]] -= 128\n",
" rgb = rgb.dot(xform.T)\n",
" np.putmask(rgb, rgb > 255, 255)\n",
" np.putmask(rgb, rgb < 0, 0)\n",
" return np.uint8(rgb)\n",
"\n",
"# Convert data and sample to YCbCr\n",
"ycbcr = rgb2ycbcr(data)\n",
"ysub = rgb2ycbcr(sub)\n",
"\n",
"# Calculate mean components\n",
"yc = list(np.mean(ysub[:,:,i]) for i in range(3))\n",
"\n",
"# Center cb and cr components of image based on sample\n",
"for i in range(1,3):\n",
" ycbcr[:,:,i] = np.clip(ycbcr[:,:,i] + (128-yc[i]), 0, 255)\n",
"\n",
"rgb = ycbcr2rgb(ycbcr) # Convert back\n",
"Image.fromarray(rgb)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.