diff --git a/scripts/check_gamma_correction.py b/scripts/check_gamma_correction.py
index bb3ce65..b1abc8f 100755
--- a/scripts/check_gamma_correction.py
+++ b/scripts/check_gamma_correction.py
@@ -31,8 +31,13 @@ imgs = np.array(imgs)
 print(imgs.shape)
 print(imgs.dtype)
 
-inv_resp = np.arange(256, dtype=np.float64)
-inv_resp[250:] = -1.0 # Use negative numbers to detect oversaturation 
+
+
+num_pixels_by_intensity = np.bincount(imgs.flat)
+print('num_pixels_by_intensity', num_pixels_by_intensity)
+
+inv_resp = np.arange(num_pixels_by_intensity.shape[0], dtype=np.float64)
+inv_resp[-1] = -1.0 # Use negative numbers to detect saturation 
 
 
 def opt_irradiance():
@@ -43,7 +48,9 @@ def opt_irradiance():
     corrected_imgs[corrected_imgs < 0] = 0
 
     denom = np.sum(times, axis=0)
-    irr = np.sum(corrected_imgs, axis=0) / denom
+    idx = (denom != 0)
+    irr = np.sum(corrected_imgs, axis=0)
+    irr[idx] /= denom[idx]
     irr[denom == 0] = -1.0
     return irr
 
@@ -51,14 +58,16 @@ def opt_inv_resp():
     generated_imgs = irradiance[np.newaxis, :, :] * exposures[:, np.newaxis, np.newaxis]
     
     num_pixels_by_intensity = np.bincount(imgs.flat, generated_imgs.flat >= 0)
+    
+    generated_imgs[generated_imgs < 0] = 0
     sum_by_intensity = np.bincount(imgs.flat, generated_imgs.flat)
     
     new_inv_resp = inv_resp
 
     idx = np.nonzero(num_pixels_by_intensity > 0)
     new_inv_resp[idx] = sum_by_intensity[idx] / num_pixels_by_intensity[idx]
-    new_inv_resp[250:] = -1.0
-    return new_inv_resp
+    new_inv_resp[-1] = -1.0 # Use negative numbers to detect saturation 
+    return new_inv_resp 
 
 def print_error():
     generated_imgs = irradiance[np.newaxis, :, :] * exposures[:, np.newaxis, np.newaxis]
@@ -66,7 +75,7 @@ def print_error():
     generated_imgs[imgs == 255] = 0
     print(np.sum(generated_imgs**2))
 
-for iter in range(3):
+for iter in range(5):
     irradiance = opt_irradiance()
     print_error()
     inv_resp = opt_inv_resp()
@@ -75,12 +84,14 @@ for iter in range(3):
 
 
 plt.figure()
-plt.plot(inv_resp)
-plt.ylabel('Img Mean')
-plt.xlabel('Exposure')
+plt.plot(inv_resp[:-1])
+plt.ylabel('Irradiance Value')
+plt.xlabel('Image Intensity')
+plt.title('Inverse Responce Function')
 
 plt.figure()
 plt.imshow(irradiance)
+plt.title('Irradiance Image')
 plt.show()