seathru_neham_hainh.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "CZPluMxpqZFP"
   },
   "outputs": [],
   "source": [
    "import collections\n",
    "import sys\n",
    "import argparse\n",
    "import numpy as np\n",
    "import sklearn as sk\n",
    "import scipy as sp\n",
    "import scipy.optimize\n",
    "import glob\n",
    "import scipy.stats\n",
    "import imageio\n",
    "import math\n",
    "from camera_pipeline_wb import after_wb\n",
    "from PIL import Image\n",
    "import torch\n",
    "import pandas as pd\n",
    "from datetime import datetime\n",
    "import rawpy\n",
    "import matplotlib\n",
    "from tqdm.notebook import tqdm\n",
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "from skimage import exposure\n",
    "from skimage.restoration import denoise_bilateral, denoise_tv_chambolle, estimate_sigma\n",
    "from skimage.morphology import closing, opening, erosion, dilation, disk, diamond, square\n",
    "torch.set_default_dtype(torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "gjy--cGksDqJ"
   },
   "outputs": [],
   "source": [
    "def process_image(filepath):\n",
    "  ''' Gets image from filepath.performs linearization and demosaicing'''\n",
    "  def get_channels(bayer_image):\n",
    "    blue = bayer_image[1::2,1::2]\n",
    "    red = bayer_image[0::2,0::2]\n",
    "    green1 = bayer_image[1::2,0::2]\n",
    "    green2 = bayer_image[0::2,1::2]\n",
    "    green = (green1 + green2)/2\n",
    "    return red, green, blue\n",
    "  \n",
    "  #Loading the RGB image\n",
    "  with rawpy.imread(filepath) as raw1:\n",
    "\n",
    "      imageio.imwrite(\"og_image.jpg\",raw1.postprocess())\n",
    "      rgb=raw1.raw_image_visible\n",
    "      saturation=raw1.camera_white_level_per_channel[0]\n",
    "      black=min(raw1.black_level_per_channel[0],rgb.min())\n",
    "      rgb=rgb.astype(np.float32)\n",
    "      rgb_linearized=(rgb-black)/(saturation-black)\n",
    "      rgb1=get_channels(rgb_linearized)\n",
    "      rgb_image=np.dstack(rgb1)\n",
    "      raw1.close()\n",
    " \n",
    "  return np.array(rgb_image)\n",
    "\n",
    "\n",
    "def image_resize(image, width = None, height = None, inter = cv2.INTER_AREA):\n",
    "\n",
    "  dim = None\n",
    "  (h, w) = image.shape[:2]  \n",
    "  if width is None and height is None:\n",
    "      return image\n",
    "  if width is not None and height is not None:\n",
    "    return cv2.resize(image,(width,height),interpolation=inter)\n",
    "  if width is None:\n",
    "      r = height / float(h)\n",
    "      dim = (int(w * r), height)\n",
    "  else:\n",
    "      r = width / float(w)\n",
    "      dim = (width, int(h * r))\n",
    "\n",
    "  resized = cv2.resize(image, dim, interpolation = inter)\n",
    "  return resized\n",
    "\n",
    "\n",
    "\n",
    "def scale(img):\n",
    "    return (img - np.min(img)) / (np.max(img) - np.min(img))\n",
    "\n",
    "def preprocess_sfm_depth_map(depths, min_depth, max_depth):\n",
    "    z_min = np.min(depths) + (min_depth * (np.max(depths) - np.min(depths)))\n",
    "    z_max = np.min(depths) + (max_depth * (np.max(depths) - np.min(depths)))\n",
    "    if max_depth != 0:\n",
    "        depths[depths == 0] = z_max\n",
    "\n",
    "    return depths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "Q6E7jqmwqyfF"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Backscatter and Preprocessing Path\n",
    "\"\"\"\n",
    "size = 1024;f=2;l=0.5;p=0.002;min_depth=0.1;max_depth=1.0;fraction = 0.01\n",
    "\n",
    "def load_image_and_depth_map(img_fname, depths_fname, size_limit = 512):\n",
    "    \n",
    "    depths = np.array(Image.open(depths_fname))\n",
    "    img  = process_image(img_fname)\n",
    "    resized_img = image_resize(img,width=size_limit)\n",
    "    depth_resized = image_resize(depths,height=resized_img.shape[0],width=size_limit)\n",
    "    return resized_img, depth_resized\n",
    "\n",
    "\n",
    "def find_backscatter_estimation_points(img, depths, num_bins=10, fraction=0.01, max_vals=100, min_depth_percent=0.0):\n",
    "    z_max, z_min = np.max(depths), np.min(depths)\n",
    "    min_depth = z_min + (min_depth_percent * (z_max - z_min))\n",
    "    z_ranges = np.linspace(z_min, z_max, num_bins + 1)\n",
    "    img_norms = np.mean(img, axis=2)\n",
    "    #img_norms = 0.2989 * img[:,:,0] + 0.5870 * img[:,:,1] + 0.1140 * img[:,:,2] #Change here\n",
    "    points_r = []\n",
    "    points_g = []\n",
    "    points_b = []\n",
    "    for i in range(len(z_ranges) - 1):\n",
    "        a, b = z_ranges[i], z_ranges[i+1]\n",
    "        locs = np.where(np.logical_and(depths > min_depth, np.logical_and(depths >= a, depths <= b)))\n",
    "        norms_in_range, px_in_range, depths_in_range = img_norms[locs], img[locs], depths[locs]\n",
    "        arr = sorted(zip(norms_in_range, px_in_range, depths_in_range), key=lambda x: x[0])\n",
    "        points = arr[:min(math.ceil(fraction * len(arr)), max_vals)]\n",
    "        points_r.extend([(z, p[0]) for n, p, z in points])\n",
    "        points_g.extend([(z, p[1]) for n, p, z in points])\n",
    "        points_b.extend([(z, p[2]) for n, p, z in points])\n",
    "\n",
    "    return np.array(points_r), np.array(points_g), np.array(points_b)\n",
    "\n",
    "def find_backscatter_values(B_pts, depths, restarts=100, max_mean_loss_fraction=0.1):\n",
    "    B_vals, B_depths = B_pts[:, 1], B_pts[:, 0]\n",
    "    z_max, z_min = np.max(depths), np.min(depths)\n",
    "    max_mean_loss = max_mean_loss_fraction * (z_max - z_min)\n",
    "    coefs = None\n",
    "    best_loss = np.inf\n",
    "\n",
    "    def estimate(depths, B_inf, beta_B, J_prime, beta_D_prime):\n",
    "        val = (B_inf * (1 - np.exp(-1 * beta_B * depths))) + (J_prime * np.exp(-1 * beta_D_prime * depths))\n",
    "        return val\n",
    "    def loss(B_inf, beta_B, J_prime, beta_D_prime):\n",
    "        val = np.mean(np.abs(B_vals - estimate(B_depths, B_inf, beta_B, J_prime, beta_D_prime)))\n",
    "        return val\n",
    "    \n",
    "    bounds_lower = [0,0,0,0]\n",
    "    bounds_upper = [1,5,1,5]\n",
    "    \n",
    "    for _ in range(restarts):\n",
    "        try:\n",
    "            optp, pcov = sp.optimize.curve_fit(\n",
    "                f=estimate,\n",
    "                xdata=B_depths,\n",
    "                ydata=B_vals,\n",
    "                p0=np.random.random(4) * bounds_upper,\n",
    "                bounds=(bounds_lower, bounds_upper),\n",
    "            )\n",
    "            l = loss(*optp)\n",
    "            if l < best_loss:\n",
    "                best_loss = l\n",
    "                coefs = optp\n",
    "        except RuntimeError as re:\n",
    "            print(re, file=sys.stderr)\n",
    "    if best_loss > max_mean_loss:\n",
    "        print('Warning: could not find accurate reconstruction. Switching to linear model.', flush=True)\n",
    "        slope, intercept, r_value, p_value, std_err = sp.stats.linregress(B_depths, B_vals)\n",
    "        BD = (slope * depths) + intercept\n",
    "        return BD, np.array([slope, intercept])\n",
    "\n",
    "    print(best_loss)\n",
    "    return estimate(depths, *coefs), coefs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "yHinslvQvGvH"
   },
   "outputs": [],
   "source": [
    "def estimate_wideband_attentuation(depths, illum, radius = 6, max_val = 10.0):\n",
    "      eps = 1E-8\n",
    "      BD = np.minimum(max_val, -np.log(illum + eps) / (np.maximum(0, depths) + eps))\n",
    "      mask = np.where(np.logical_and(depths > eps, illum > eps), 1, 0)\n",
    "      refined_attenuations = denoise_bilateral(closing(np.maximum(0, BD * mask), disk(radius)))\n",
    "      return refined_attenuations, []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "ktRTnWbXvPaz"
   },
   "outputs": [],
   "source": [
    "def filter_data(X, Y, radius_fraction=0.01):\n",
    "    idxs = np.argsort(X)\n",
    "    print(idxs)\n",
    "    idxs = idxs.astype(np.int32)\n",
    "    X_s = X[idxs]\n",
    "    idxs = idxs.astype(np.int32)\n",
    "    Y_s = Y[idxs]\n",
    "    x_max, x_min = np.max(X), np.min(X)\n",
    "    radius = (radius_fraction * (x_max - x_min))\n",
    "    ds = np.cumsum(X_s - np.roll(X_s, (1,)))\n",
    "    dX = [X_s[0]]\n",
    "    dY = [Y_s[0]]\n",
    "    tempX = []\n",
    "    tempY = []\n",
    "    pos = 0\n",
    "    for i in range(1, ds.shape[0]):\n",
    "        if ds[i] - ds[pos] >= radius:\n",
    "            tempX.append(X_s[i])\n",
    "            tempY.append(Y_s[i])\n",
    "            idxs = np.argsort(tempY)\n",
    "            med_idx = len(idxs) // 2\n",
    "            dX.append(tempX[med_idx])\n",
    "            dY.append(tempY[med_idx])\n",
    "            pos = i\n",
    "        else:\n",
    "            tempX.append(X_s[i])\n",
    "            tempY.append(Y_s[i])\n",
    "    return np.array(dX), np.array(dY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "Ee016bCKusOM"
   },
   "outputs": [],
   "source": [
    "\n",
    "'''\n",
    "\n",
    "Estimate coefficients for the 2-term exponential\n",
    "describing the wideband attenuation\n",
    "'''\n",
    "def refine_wideband_attentuation(depths, illum, estimation, restarts=25, min_depth_fraction = 0.1, max_mean_loss_fraction=np.inf, l=1.0, radius_fraction=0.01):\n",
    "    eps = 1E-8\n",
    "    z_max, z_min = np.max(depths), np.min(depths)\n",
    "    min_depth = z_min + (min_depth_fraction * (z_max - z_min))\n",
    "    max_mean_loss = max_mean_loss_fraction * (z_max - z_min)\n",
    "    coefs = None\n",
    "    best_loss = np.inf\n",
    "    locs = np.where(np.logical_and(illum > 0, np.logical_and(depths > min_depth, estimation > eps)))\n",
    "    def calculate_reconstructed_depths(depths, illum, a, b, c, d):\n",
    "        eps = 1E-5\n",
    "        res = -np.log(illum + eps) / (calculate_beta_D(depths, a, b, c, d) + eps)\n",
    "        return res\n",
    "    def loss(a, b, c, d):\n",
    "        return np.mean(np.abs(depths[locs] - calculate_reconstructed_depths(depths[locs], illum[locs], a, b, c, d)))\n",
    "    dX, dY = filter_data(depths[locs], estimation[locs], radius_fraction)\n",
    "    for _ in range(restarts):\n",
    "        try:\n",
    "            optp, pcov = sp.optimize.curve_fit(\n",
    "                f=calculate_beta_D,\n",
    "                xdata=dX,\n",
    "                ydata=dY,\n",
    "                p0=np.abs(np.random.random(4)) * np.array([1., -1., 1., -1.]),\n",
    "                bounds=([0, -100, 0, -100], [100, 0, 100, 0]))\n",
    "            L = loss(*optp)\n",
    "            if L < best_loss:\n",
    "                best_loss = L\n",
    "                coefs = optp\n",
    "        except RuntimeError as re:\n",
    "            print(re, file=sys.stderr)\n",
    "    # Uncomment to see the regression\n",
    "    # plt.clf()\n",
    "    # plt.scatter(depths[locs], estimation[locs])\n",
    "    # plt.plot(np.sort(depths[locs]), calculate_beta_D(np.sort(depths[locs]), *coefs))\n",
    "    # plt.show()\n",
    "    print(f'Found best loss {best_loss}', flush=True)\n",
    "    BD = l * calculate_beta_D(depths, *coefs)\n",
    "    return BD, coefs,best_loss\n",
    "\n",
    "  \n",
    "def wbalance_no_red_10p(img):\n",
    "    dg = 1.0 / np.mean(np.sort(img[:, :, 1], axis=None)[int(round(-1 * np.size(img[:, :, 0]) * 0.1)):])\n",
    "    db = 1.0 / np.mean(np.sort(img[:, :, 2], axis=None)[int(round(-1 * np.size(img[:, :, 0]) * 0.1)):])\n",
    "    dsum = dg + db\n",
    "    dg = dg / dsum * 2.\n",
    "    db = db / dsum * 2.\n",
    "    img[:, :, 0] *= (db + dg) / 2\n",
    "    img[:, :, 1] *= dg\n",
    "    img[:, :, 2] *= db\n",
    "    return img\n",
    "    \n",
    "def recover_image(img, depths, B, beta_D):\n",
    "    \n",
    "    res = (img - B) * np.exp(beta_D * np.expand_dims(depths, axis=2))\n",
    "    print(res.max(),res.min())\n",
    "    print(np.percentile(res, range(0,100,5)))\n",
    "    res = np.maximum(0.0, np.minimum(1.0, res))\n",
    "    res = scale(wbalance_no_red_10p(res))\n",
    "    return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "E23EzQQUESlK"
   },
   "outputs": [],
   "source": [
    "'''\n",
    "Estimate coefficients for the 2-term exponential\n",
    "describing the wideband attenuation\n",
    "'''\n",
    "'''\n",
    "Calculate the values of beta_D for an image from the depths, illuminations, and constants\n",
    "'''\n",
    "def calculate_beta_D(depths, a,b,c,d):\n",
    "    return a * np.exp(-b * depths)+ c * np.exp(-d * depths)\n",
    "\n",
    "\n",
    "def refine_wideband_attentuation_neham(depths, illum, estimation, restarts=10, min_depth_fraction = 0.1, max_mean_loss_fraction=np.inf, l=1.0, radius_fraction=0.01):\n",
    "    eps = 1E-8\n",
    "    min_depth = np.percentile(depths, 10, axis=None, out=None) \n",
    "    max_depth = np.percentile(depths, 90, axis=None, out=None) \n",
    "    best_loss = np.inf\n",
    "\n",
    "    pts_to_consider = np.where(np.logical_and(depths>=min_depth,depths<=max_depth))\n",
    "    #print(pts_to_consider)\n",
    "    dX = depths[pts_to_consider]\n",
    "    dY = depths[pts_to_consider]\n",
    "\n",
    "    def calculate_reconstructed_depths(depths, illum, a, b,c,d):\n",
    "        eps = 1E-5\n",
    "        res = -np.log(illum + eps) / (calculate_beta_D(depths, a, b,c,d) + eps)\n",
    "        return res\n",
    "    \n",
    "    def loss(a, b,c,d):\n",
    "        return np.mean(np.abs(depths[pts_to_consider] - calculate_reconstructed_depths(depths[pts_to_consider], illum[pts_to_consider], a, b,c,d)))\n",
    "    \n",
    "    \n",
    "    for _ in range(restarts):\n",
    "        try:\n",
    "            optp, pcov = sp.optimize.curve_fit(\n",
    "                f=calculate_beta_D,\n",
    "                xdata=dX,\n",
    "                ydata=dY,\n",
    "                p0=np.abs(np.random.random(4)) * np.array([5., 3.,4,5]),\n",
    "                bounds=([0, 0,0,0], [100, 100,100,100]))\n",
    "            L = loss(*optp)\n",
    "            if L < best_loss:\n",
    "                best_loss = L\n",
    "                coefs = optp\n",
    "        except RuntimeError as re:\n",
    "            print(re, file=sys.stderr)\n",
    "    # Uncomment to see the regression\n",
    "    # plt.clf()\n",
    "    # plt.scatter(depths[locs], estimation[locs])\n",
    "    # plt.plot(np.sort(depths[locs]), calculate_beta_D(np.sort(depths[locs]), *coefs))\n",
    "    # plt.show()\n",
    "    print(f'Found best loss {best_loss}', flush=True)\n",
    "    BD = l * calculate_beta_D(depths, *coefs)\n",
    "    return BD, coefs,best_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "faJ8prTb_Ezy"
   },
   "outputs": [],
   "source": [
    "def depth_iterative(without_backscatter,pd,normalized_depth,count,epochs):\n",
    "  dc=torch.zeros((without_backscatter.shape[0],without_backscatter.shape[1])).cuda()\n",
    "  pd=torch.tensor(pd, dtype=torch.float32).cuda()\n",
    "  normalized_depth=torch.from_numpy(normalized_depth).cuda()\n",
    "  count=torch.from_numpy(count).cuda()\n",
    "\n",
    "  for i in tqdm(range(epochs)):\n",
    "\n",
    "    dc_temp=torch.zeros((without_backscatter.shape[0],without_backscatter.shape[1])).cuda()\n",
    "    #Left Condition\n",
    "    dc_temp[:,1:] = dc_temp[:,1:] + dc[:,:-1]\n",
    "    #Right Condition\n",
    "    dc_temp[:,:-1] = dc_temp[:,:-1] + dc[:,1:]\n",
    "    #Top Condition\n",
    "    dc_temp[1:,:] = dc_temp[1:,:] + dc[:-1,:]\n",
    "    #Bottom Condition\n",
    "    dc_temp[:-1,:] = dc_temp[:-1,:] + dc[1:,:]\n",
    "\n",
    "    dc_temp=torch.div(dc_temp,count)\n",
    "\n",
    "    dc = pd*normalized_depth + (1-pd)* dc_temp\n",
    "\n",
    "    if i%5000==0:\n",
    "      print(\"Epoch: \",i,\"Mean: \",dc.mean())\n",
    "    \n",
    "  return dc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "z_CJGSTgHSm4"
   },
   "outputs": [],
   "source": [
    "def determine_depth_threshold(without_backscatter,normalized_depth,depth_image,pd,epochs=25000):\n",
    "  \n",
    "  count=4*np.ones(((without_backscatter.shape[0],without_backscatter.shape[1])))\n",
    "  count[0,:]=count[0,:]-1\n",
    "  count[-1,:]=count[-1,:]-1\n",
    "  count[:,0]=count[:,0]-1\n",
    "  count[:,-1]=count[:,-1  ]-1\n",
    "  count=count.astype(np.float32)\n",
    "  dc=depth_iterative(without_backscatter,pd,normalized_depth,count,epochs)\n",
    "  return dc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "2yNfNHx6w1S5"
   },
   "outputs": [],
   "source": [
    "def iterative_func(count,without_backscatter,lc,rc,tc,bc,epochs,p):\n",
    "  p=torch.tensor(p, dtype=torch.float32).cuda()\n",
    "  lc=torch.from_numpy(lc).cuda()\n",
    "  rc=torch.from_numpy(rc).cuda()\n",
    "  tc=torch.from_numpy(tc).cuda()\n",
    "  bc=torch.from_numpy(bc).cuda()\n",
    "  count=torch.from_numpy(count).cuda()\n",
    "\n",
    "  ac=torch.zeros((without_backscatter.shape[0],without_backscatter.shape[1],without_backscatter.shape[2])).cuda()\n",
    "  without_backscatter=torch.from_numpy(without_backscatter).cuda()\n",
    "  for i in tqdm(range(epochs)):\n",
    "    #temp1=compute_ac_dash(temp,lc,rc,tc,bc,ac)\n",
    "    #ac_temp = np.divide(compute_ac_dash(temp,lc,rc,tc,bc,ac), count, out=np.zeros_like(count), where=count!=0)\n",
    "    ac_temp=torch.zeros((without_backscatter.shape[0],without_backscatter.shape[1],without_backscatter.shape[2])).cuda()\n",
    "    #Left Condition\n",
    "    ac_temp[:,1:] = ac_temp[:,1:] + torch.mul(lc[:,1:,:],ac[:,:-1,:])    \n",
    "\n",
    "    #Right Condition\n",
    "    ac_temp[:,:-1] = ac_temp[:,:-1] + torch.mul(rc[:,:-1],ac[:,1:])\n",
    "    #Top Condition\n",
    "    ac_temp[1:,:] = ac_temp[1:,:] + torch.mul(tc[1:,:],ac[:-1,:])\n",
    "    #Bottom Condition\n",
    "    ac_temp[:-1,:] = ac_temp[:-1,:] + torch.mul(bc[:-1,:],ac[1:,:])\n",
    "    \n",
    "    ac_temp=torch.div(ac_temp,count)\n",
    "\n",
    "    ac_temp[torch.isnan(ac_temp)] = 0 \n",
    "       \n",
    "    ac = p*without_backscatter + (1-p)* ac_temp\n",
    "    if i%5000==0:\n",
    "      print(\"Epoch: \",i,\"Mean: \",ac.mean())\n",
    "    \n",
    "  \n",
    "  return ac\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "_IKaS2zHHVt0"
   },
   "outputs": [],
   "source": [
    "def LSAC(without_backscatter,normalized_depth,threshold,p,epochs=25000):\n",
    "  \n",
    "  def determine_neighbourhood(depth_image,threshold):\n",
    "    #Check Left Condition \n",
    "    lc=abs(depth_image[:,1:]-depth_image[:,:-1])<threshold[:,1:]\n",
    "    lc=np.hstack((np.zeros((lc.shape[0],1)),lc))\n",
    "    lc=lc[:,:,np.newaxis]\n",
    "    #Check Right Condition \n",
    "    rc=abs(depth_image[:,:-1]-depth_image[:,1:])<threshold[:,:-1]\n",
    "    rc=np.hstack((rc,np.zeros((rc.shape[0],1))))\n",
    "    rc=rc[:,:,np.newaxis]\n",
    "    \n",
    "    #Check Top Condition \n",
    "    tc=abs(depth_image[1:,:]-depth_image[:-1,:])<threshold[1:,:]\n",
    "    tc=np.vstack((np.zeros((1,tc.shape[1])),tc))\n",
    "    tc=tc[:,:,np.newaxis]\n",
    "\n",
    "    #Check Bottom Condition \n",
    "    bc=abs(depth_image[:-1,:]-depth_image[1:,:])<threshold[:-1,:]\n",
    "    bc=np.vstack((bc,np.zeros((1,bc.shape[1]))))\n",
    "    bc=bc[:,:,np.newaxis]\n",
    "    lc=lc.astype(np.float32)\n",
    "    rc=rc.astype(np.float32)\n",
    "    tc=tc.astype(np.float32)\n",
    "    bc=bc.astype(np.float32)\n",
    "\n",
    "    return lc,rc,tc,bc\n",
    "  \n",
    "\n",
    "  lc,rc,tc,bc=determine_neighbourhood(normalized_depth,threshold=threshold)\n",
    "  count=lc+rc+tc+bc\n",
    "  without_backscatter=without_backscatter.astype(np.float32)\n",
    "  ac=iterative_func(count,without_backscatter,lc,rc,tc,bc,epochs,p)    \n",
    "  return ac\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "coz2_iKUqpGu",
    "outputId": "e95457e9-f9a8-44cc-922d-141a622b0c5c"
   },
   "outputs": [],
   "source": [
    "size = 1024;f=2;l=0.5;p=0.002;min_depth=0.1;max_depth=1.0;fraction = 0.01\n",
    "\n",
    "def seathru(img_path,depth_path,save_path):\n",
    "\n",
    "    raw = rawpy.imread(img_path)\n",
    "    img, depths = load_image_and_depth_map(img_path,depth_path)\n",
    "    save_coeffs = {}\n",
    "    depths = preprocess_sfm_depth_map(depths,min_depth,max_depth)\n",
    "    print('Estimating backscatter...', flush=True)\n",
    "    ptsR, ptsG, ptsB = find_backscatter_estimation_points(img, depths, fraction=0.01, min_depth_percent=min_depth)\n",
    "\n",
    "    print('Finding backscatter coefficients...', flush=True)\n",
    "    Br, coefsR = find_backscatter_values(ptsR, depths, restarts=100)\n",
    "    Bg, coefsG = find_backscatter_values(ptsG, depths, restarts=100)\n",
    "    Bb, coefsB = find_backscatter_values(ptsB, depths, restarts=100)\n",
    "    B = np.stack([Br, Bg, Bb], axis=2)\n",
    "    \n",
    "    without_backscatter = np.clip(img-B,0,1)\n",
    "    print('Estimating wideband attenuation...', flush=True)\n",
    "\n",
    "    pc=5e-4\n",
    "    pd=5e-4\n",
    "\n",
    "    normalized_depth = (depths-np.min(depths))/(np.max(depths)-np.min(depths))\n",
    "\n",
    "    dc=determine_depth_threshold(without_backscatter.copy(),normalized_depth.copy(),depths.copy(),pd,epochs=25000)\n",
    "    dc=dc.cpu().detach().numpy()  \n",
    "\n",
    "    #Part 4\n",
    "    threshold=0.1*dc\n",
    "\n",
    "    ac=LSAC(without_backscatter,normalized_depth,threshold,pd,epochs=25000)\n",
    "    ac=ac.cpu().detach().numpy()  \n",
    "    illumination=2*ac\n",
    "    illumination[np.where(illumination==0)] = illumination[np.where(illumination==0)]+1e-5\n",
    "\n",
    "    illR = illumination[:,:,0]\n",
    "    illG = illumination[:,:,1]\n",
    "    illB = illumination[:,:,2]\n",
    "\n",
    "    beta_D_r, _ = estimate_wideband_attentuation(depths, illR)\n",
    "    #refined_beta_D_r, coefsR,best_loss_R = refine_wideband_attentuation_neham(depths, illR, beta_D_r, radius_fraction = fraction, l=l)\n",
    "\n",
    "    beta_D_g, _ = estimate_wideband_attentuation(depths, illG)\n",
    "    #refined_beta_D_g, coefsG,best_loss_G = refine_wideband_attentuation_neham(depths, illG, beta_D_g, radius_fraction = fraction, l=l)\n",
    "\n",
    "    beta_D_b, _ = estimate_wideband_attentuation(depths, illB)\n",
    "    #refined_beta_D_b, coefsB,best_loss_B = refine_wideband_attentuation_neham(depths, illB, beta_D_b, radius_fraction = fraction, l=l)\n",
    "\n",
    "    save_coeffs[\"file_name\"] = img_path \n",
    "    save_coeffs[\"depth_mean\"] = depths.mean()\n",
    "\n",
    "    print('Reconstructing image...', flush=True)\n",
    "    B = np.stack([Br, Bg, Bb], axis=2)\n",
    "    beta_D_og = np.stack([beta_D_r,beta_D_g,beta_D_b], axis=2)\n",
    "    #refined_beta_D_og = np.stack([refined_beta_D_r,refined_beta_D_g,refined_beta_D_b], axis=2)\n",
    "\n",
    "    recovered_og = recover_image(img, depths, B, beta_D_og)\n",
    "\n",
    "    #recovered_refined = recover_image(img, depths, B, refined_beta_D_og)\n",
    "    \n",
    "    imageio.imwrite(save_path,recovered_og)\n",
    "    imageio.imwrite(\"results/\"+\"postprocessed\"+save_path.split('/')[-1],raw.postprocess(half_size=True))\n",
    "    after_wb(raw,img,\"results/\"+\"original\"+save_path.split('/')[-1])\n",
    "    return np.exp(beta_D_og * np.expand_dims(depths, axis=2)),depths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "QKG1qmZ4QXXo"
   },
   "outputs": [],
   "source": [
    "def beta_D(coefs,depths):\n",
    "  return coefs[0]*np.exp(-coefs[1]*depths) + coefs[2]*np.exp(-coefs[3]*depths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 282
    },
    "id": "Tw9fiAJwHP65",
    "outputId": "c01942bb-defe-4010-e8a7-b6dce689078d"
   },
   "outputs": [],
   "source": [
    "files = glob.glob(\"/home/neham/uw_datasets/SeaThru_Final/*\")\n",
    "input_files = []\n",
    "input_depths = []\n",
    "outputs = []\n",
    "count = 0\n",
    "for i in files:\n",
    "    raw_files = glob.glob(i+\"/Raw/*\")[-2:]\n",
    "    list_depth = glob.glob(i+\"/depthMaps/*\")\n",
    "    for file_name in raw_files:\n",
    "        depth_num = file_name.split(\".\")[0][-3:]\n",
    "        depth_file = [i for i in list_depth if depth_num in i]\n",
    "        if len(depth_file)==0:\n",
    "          continue\n",
    "        file_d = depth_file[0]\n",
    "        output_name = \"results/\" + file_name.split(\"/\")[-1][:-3] + \"png\"\n",
    "        count+=1\n",
    "        input_files.append(file_name)\n",
    "        input_depths.append(file_d)\n",
    "        outputs.append(output_name)\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/home/neham/uw_datasets/SeaThru_Final/SeaThru_8/Raw/T_S04863.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_8/Raw/T_S04895.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_4/Raw/T_S03507.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_4/Raw/T_S03488.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_11/Raw/LFT_3414.NEF',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_11/Raw/LFT_3413.NEF',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_3/Raw/T_S03386.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_3/Raw/T_S03422.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_0/Raw/T_S03132.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_0/Raw/T_S03107.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_9/Raw/RGT_0217.NEF',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_9/Raw/RGT_0178.NEF',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_2/Raw/T_S03336.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_2/Raw/T_S03282.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_1/Raw/T_S03217.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_1/Raw/T_S03235.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_10/Raw/RGT_0294.NEF',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_10/Raw/RGT_0268.NEF',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_5/Raw/T_S03673.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_5/Raw/T_S03585.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_7/Raw/T_S03852.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_7/Raw/T_S03812.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_6/Raw/T_S03773.ARW',\n",
       " '/home/neham/uw_datasets/SeaThru_Final/SeaThru_6/Raw/T_S03745.ARW']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "393e99b9770a4ea5ad1646c8ec7ee3c3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Estimating backscatter...\n",
      "Finding backscatter coefficients...\n",
      "0.0010819504\n",
      "0.0050101727\n",
      "0.0033833594\n",
      "Estimating wideband attenuation...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f00f901309704b3bae44de072cb2b8b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  0 Mean:  tensor(0.0003, device='cuda:0')\n",
      "Epoch:  5000 Mean:  tensor(0.5438, device='cuda:0')\n",
      "Epoch:  10000 Mean:  tensor(0.5884, device='cuda:0')\n",
      "Epoch:  15000 Mean:  tensor(0.5920, device='cuda:0')\n",
      "Epoch:  20000 Mean:  tensor(0.5923, device='cuda:0')\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e91705ffd00a4377ba20ec1de55b569e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  0 Mean:  tensor(2.2894e-05, device='cuda:0')\n",
      "Epoch:  5000 Mean:  tensor(0.0418, device='cuda:0')\n",
      "Epoch:  10000 Mean:  tensor(0.0452, device='cuda:0')\n",
      "Epoch:  15000 Mean:  tensor(0.0455, device='cuda:0')\n",
      "Epoch:  20000 Mean:  tensor(0.0455, device='cuda:0')\n",
      "Reconstructing image...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Lossy conversion from float64 to uint8. Range [0, 1]. Convert image to uint8 prior to saving to suppress this warning.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "596372043037.5481 -6675015257.106591\n",
      "[-6.67501526e+09  1.65949170e-02  8.13733900e-02  1.36490499e-01\n",
      "  1.94589236e-01  2.59015039e-01  3.24243811e-01  3.83502841e-01\n",
      "  4.31728178e-01  4.66206202e-01  4.86514227e-01  4.98668791e-01\n",
      "  5.07333012e-01  5.17404253e-01  5.45407583e-01  6.03395343e-01\n",
      "  6.71001404e-01  7.49311505e-01  8.49061672e-01  1.02834293e+00]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Lossy conversion from float64 to uint8. Range [0, 1]. Convert image to uint8 prior to saving to suppress this warning.\n"
     ]
    }
   ],
   "source": [
    "for i,j,k in tqdm(zip(input_files,input_depths,outputs)):\n",
    "    try:\n",
    "        direct_attenuation,depths = seathru(i,j,k)\n",
    "        break\n",
    "    except:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.9585354 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_8/depthMaps/depthT_S04863.tif\n",
      "2.8167076 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_8/depthMaps/depthT_S04895.tif\n",
      "0.97354454 0.4955313 /home/neham/uw_datasets/SeaThru_Final/SeaThru_4/depthMaps/depthT_S03507.tif\n",
      "1.2640331 0.6778563 /home/neham/uw_datasets/SeaThru_Final/SeaThru_4/depthMaps/depthT_S03488.tif\n",
      "4.4719105 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_11/depthMaps/depthLFT_3414.tif\n",
      "5.2783165 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_11/depthMaps/depthLFT_3413.tif\n",
      "1.2959849 0.6499919 /home/neham/uw_datasets/SeaThru_Final/SeaThru_3/depthMaps/depthT_S03386.tif\n",
      "1.1321101 0.6963256 /home/neham/uw_datasets/SeaThru_Final/SeaThru_3/depthMaps/depthT_S03422.tif\n",
      "1.817156 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_0/depthMaps/depthT_S03132.tiff\n",
      "2.0637763 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_0/depthMaps/depthT_S03107.tiff\n",
      "2.998172 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_9/depthMaps/depthRGT_0217.tif\n",
      "4.340022 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_9/depthMaps/depthRGT_0178.tif\n",
      "1.2038705 0.5641684 /home/neham/uw_datasets/SeaThru_Final/SeaThru_2/depthMaps/depthT_S03336.tif\n",
      "1.5651559 0.6213467 /home/neham/uw_datasets/SeaThru_Final/SeaThru_2/depthMaps/depthT_S03282.tif\n",
      "1.4165417 0.63216543 /home/neham/uw_datasets/SeaThru_Final/SeaThru_1/depthMaps/depthT_S03217.tif\n",
      "1.5609033 0.6422095 /home/neham/uw_datasets/SeaThru_Final/SeaThru_1/depthMaps/depthT_S03235.tif\n",
      "4.213322 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_10/depthMaps/depthRGT_0294.tif\n",
      "4.812146 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_10/depthMaps/depthRGT_0268.tif\n",
      "4.673552 2.4521232 /home/neham/uw_datasets/SeaThru_Final/SeaThru_5/depthMaps/depthT_S03673.tif\n",
      "4.6130753 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_5/depthMaps/depthT_S03585.tif\n",
      "8.012659 0.0 /home/neham/uw_datasets/SeaThru_Final/SeaThru_7/depthMaps/depthT_S03852.tif\n",
      "5.784081 3.4143937 /home/neham/uw_datasets/SeaThru_Final/SeaThru_7/depthMaps/depthT_S03812.tif\n",
      "6.321379 3.7662284 /home/neham/uw_datasets/SeaThru_Final/SeaThru_6/depthMaps/depthT_S03773.tif\n",
      "7.1645703 3.5292609 /home/neham/uw_datasets/SeaThru_Final/SeaThru_6/depthMaps/depthT_S03745.tif\n"
     ]
    }
   ],
   "source": [
    "for i in input_depths:\n",
    "    depths = np.array(Image.open(i))\n",
    "    print(depths.max(),depths.min(),i)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "seathru_neham_hainh.ipynb",
   "provenance": []
  },
  "interpreter": {
   "hash": "852a3c0ef8269d1055315cf744c6b9714e6cbc795cf4b9190293d035134ef204"
  },
  "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}