test_solar_script.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "698c8a7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from albumentations.pytorch import ToTensorV2\n",
    "from torch.utils.data import DataLoader\n",
    "import torch,numpy as np\n",
    "import torch.utils.data as data\n",
    "from glob import glob\n",
    "import albumentations as A\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "\n",
    "transform = A.Compose([  # Color normalizations used while training\n",
    "    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225),max_pixel_value=255.0, p=1.0),\n",
    "    ToTensorV2(),\n",
    "    \n",
    "])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d6af4808",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The current device is cuda:0\n"
     ]
    }
   ],
   "source": [
    "from patchify import unpatchify, patchify\n",
    "import os, cv2, torch\n",
    "from glob import glob\n",
    "import csv\n",
    "from skimage.measure import label, regionprops, regionprops_table\n",
    "import segmentation_models_pytorch as smp\n",
    "\n",
    "'''\n",
    "# Model definition - trained on part of PV dataset\n",
    "model = smp.UnetPlusPlus(\n",
    "    encoder_name=\"efficientnet-b4\",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7\n",
    "    encoder_weights=\"imagenet\",     # use `imagenet` pre-trained weights for encoder initialization\n",
    "    in_channels=3,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)\n",
    "    classes=1,                      # model output channels (number of classes in your dataset)\n",
    ")\n",
    "'''\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"The current device is {device}\")\n",
    "\n",
    "model=torch.load('upp_trained_model.pth') #load trained model\n",
    "\n",
    "model.to(device)\n",
    "model.eval()\n",
    "if not os.path.exists('out'):  #make output dirs\n",
    "    os.makedirs('out')\n",
    "\n",
    "\n",
    "for file in glob(os.path.join('./test','*.jpg')):  #for each image in test\n",
    "    basename = os.path.basename(file)\n",
    "    \n",
    "    orig_image = cv2.imread(file)  # read image\n",
    "    h = orig_image.shape[0]\n",
    "    w = orig_image.shape[1]\n",
    "    \n",
    "    \n",
    "    jgw_name = file[:-2]+'gw'     # read  jgw file\n",
    "    with open(jgw_name,'r') as f:\n",
    "        entries = [float(line.split()[0]) for line in f]\n",
    "    eastings_origin = entries[4]\n",
    "    northings_origin = entries[5]\n",
    "    eastings_step = entries[0]\n",
    "    northings_step = entries[3]\n",
    "    \n",
    "        \n",
    "    for ratio in [0.485, 1]:   # run inference at different resolutions\n",
    "        if not os.path.exists('out/'+str(ratio)):\n",
    "            os.makedirs('out/'+str(ratio))\n",
    "        \n",
    "        full_image = cv2.resize(orig_image, None, fx= ratio, fy= ratio, interpolation= cv2.INTER_LINEAR)\n",
    "\n",
    "\n",
    "        extra_h = 256 - (full_image.shape[0]%256)\n",
    "        extra_w = 256 - (full_image.shape[1]%256)\n",
    "        padded_full = cv2.copyMakeBorder(full_image.copy(),0,extra_h,0,extra_w,cv2.BORDER_CONSTANT,value=(255,0,0))\n",
    "        padded_full_rgb = cv2.cvtColor(padded_full, cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        patches = patchify(padded_full_rgb, (256,256,3), step=256) #split high res image to patches for inference\n",
    "        #patches = patches.reshape((:,:,1,256,256,3))\n",
    "        patch_results = []\n",
    "        for x_patch in patches: #run inference on each patch\n",
    "            row_patches = []\n",
    "            for xy_patch in x_patch:\n",
    "                curr_patch = np.array(xy_patch[0])\n",
    "                #print(curr_patch.shape)\n",
    "                transformed =  transform(image=curr_patch)['image'][None,:]\n",
    "                transformed = transformed.to(device)\n",
    "                #print(transformed.shape)\n",
    "                out = model(transformed)\n",
    "                outer = torch.max(out,1).indices.cpu().detach().numpy()[0]\n",
    "                row_patches.append(outer*255)\n",
    "\n",
    "\n",
    "            patch_results.append(row_patches.copy())\n",
    "\n",
    "\n",
    "        patch_results = np.array(patch_results,dtype='uint8')\n",
    "        #print(patch_results.shape,patches.shape,padded_full.shape[:-1])\n",
    "        reconstructed_image = unpatchify(patch_results, padded_full.shape[:-1])  #stitch masks of each patch\n",
    "        #print(reconstructed_image.shape)\n",
    "        \n",
    "\n",
    "        label_img = label(reconstructed_image)\n",
    "        regions = regionprops(label_img)  #segment groups as regions\n",
    "        \n",
    "        #bring to original size\n",
    "        p_out = padded_full[:full_image.shape[0],:full_image.shape[1],:]\n",
    "        p_out = cv2.resize(p_out, None, fx=1/ratio, fy= 1/ratio, interpolation= cv2.INTER_LINEAR)\n",
    "        reconstructed_image = reconstructed_image[:full_image.shape[0],:full_image.shape[1]]\n",
    "        reconstructed_image = cv2.resize(reconstructed_image, None, fx=1/ratio, fy= 1/ratio, interpolation= cv2.INTER_LINEAR)\n",
    "        \n",
    " \n",
    "        \n",
    "        bbox_bng = []\n",
    "        for i,r in enumerate(regions):  #treat each region as a detection- draw bbox in image and write to csv\n",
    "            #print(i,r.centroid,r.area)\n",
    "            if r.area<400:   #skip small detections\n",
    "                continue\n",
    "            minr, minc, maxr, maxc = r.bbox[0]/ratio, r.bbox[1]/ratio, r.bbox[2]/ratio, r.bbox[3]/ratio\n",
    "            \n",
    "            cv2.rectangle(p_out, (int(minc), int(minr)), (int(maxc),int(maxr)), (0,69, 255),2)\n",
    "            \n",
    "            \n",
    "            bbox_bng.append([i, eastings_origin+(minc*eastings_step), northings_origin+(minr*northings_step),\n",
    "                            eastings_origin+(maxc*eastings_step), northings_origin+(maxr*northings_step)])\n",
    "            # format of csv file : [id of detection,easting start, northing start, easting end, northing end ]\n",
    "\n",
    "        cv2.imwrite('out/'+str(ratio)+'/'+basename,p_out) # post-processed image\n",
    "        with open('out/'+str(ratio)+'/'+basename[:-4]+'.csv', 'w', encoding='UTF8') as f: # write bng coordinates of detections in csv\n",
    "            writer = csv.writer(f)\n",
    "            for row in bbox_bng:\n",
    "            # write a row to the csv file\n",
    "                writer.writerow(row)\n",
    "        \n",
    "        \n",
    "            \n",
    "    \n",
    "    \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f92c3a9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "283949e3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}