Crop-Segmentation-and-Yield-Count-using-Graph-based-approach.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "938eb9dd",
   "metadata": {},
   "source": [
    "# Crop Segmentation and Yield Count using Graph based methods"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9aed0e",
   "metadata": {},
   "source": [
    "### 1. Image segmentation using SLIC and Graph Cut method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f67cca8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Height:  720\n",
      "Width:   720\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3991/3732865020.py:16: FutureWarning: skimage.measure.label's indexing starts from 0. In future version it will start from 1. To disable this warning, explicitely set the `start_label` parameter to 1.\n",
      "  segments = slic(img, n_segments=500, compactness=18.5)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Height:  720\n",
      "Width:   720\n",
      "Height:  720\n",
      "Width:   720\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import sys\n",
    "import os\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage.segmentation import slic\n",
    "from skimage.segmentation import mark_boundaries\n",
    "from skimage.util import img_as_float\n",
    "import maxflow\n",
    "from scipy.spatial import Delaunay\n",
    "\n",
    "# Calculate the SLIC superpixels, their histograms and neighbors\n",
    "def superpixels_histograms_neighbors(img):\n",
    "    # SLIC\n",
    "    segments = slic(img, n_segments=500, compactness=18.5)\n",
    "    segments_ids = np.unique(segments)\n",
    "\n",
    "    # centers\n",
    "    centers = np.array([np.mean(np.nonzero(segments==i),axis=1) for i in segments_ids])\n",
    "\n",
    "    # H-S histograms for all superpixels\n",
    "    hsv = cv2.cvtColor(img.astype('float32'), cv2.COLOR_BGR2HSV)\n",
    "    bins = [20, 20] # H = S = 20\n",
    "    ranges = [0, 360, 0, 1] # H: [0, 360], S: [0, 1]\n",
    "    colors_hists = np.float32([cv2.calcHist([hsv],[0, 1], np.uint8(segments==i), bins, ranges).flatten() for i in segments_ids])\n",
    "\n",
    "    # neighbors via Delaunay tesselation\n",
    "    tri = Delaunay(centers)\n",
    "    return (centers,colors_hists,segments,tri.vertex_neighbor_vertices)\n",
    "\n",
    "# Get superpixels IDs for FG and BG from marking\n",
    "def find_superpixels_under_marking(marking, superpixels):\n",
    "    fg_segments = np.unique(superpixels[marking[:,:,0]!=255])\n",
    "    bg_segments = np.unique(superpixels[marking[:,:,2]!=255])\n",
    "    return (fg_segments, bg_segments)\n",
    "\n",
    "# Sum up the histograms for a given selection of superpixel IDs, normalize\n",
    "def cumulative_histogram_for_superpixels(ids, histograms):\n",
    "    h = np.sum(histograms[ids],axis=0)\n",
    "    return h / h.sum()\n",
    "\n",
    "# Get a bool mask of the pixels for a given selection of superpixel IDs\n",
    "def pixels_for_segment_selection(superpixels_labels, selection):\n",
    "    pixels_mask = np.where(np.isin(superpixels_labels, selection), True, False)\n",
    "    return pixels_mask\n",
    "\n",
    "# Get a normalized version of the given histograms (divide by sum)\n",
    "def normalize_histograms(histograms):\n",
    "    return np.float32([h / h.sum() for h in histograms])\n",
    "\n",
    "# Perform graph cut using superpixels histograms\n",
    "def do_graph_cut(fgbg_hists, fgbg_superpixels, norm_hists, neighbors):\n",
    "    num_nodes = norm_hists.shape[0]\n",
    "    # Create a graph of N nodes, and estimate of 5 edges per node\n",
    "    g = maxflow.Graph[float](num_nodes, num_nodes * 5)\n",
    "    # Add N nodes\n",
    "    nodes = g.add_nodes(num_nodes)\n",
    "\n",
    "    hist_comp_alg = cv2.HISTCMP_KL_DIV\n",
    "\n",
    "    # Smoothness term: cost between neighbors\n",
    "    indptr,indices = neighbors\n",
    "    for i in range(len(indptr)-1):\n",
    "        N = indices[indptr[i]:indptr[i+1]] # list of neighbor superpixels\n",
    "        hi = norm_hists[i]                 # histogram for center\n",
    "        for n in N:\n",
    "            if (n < 0) or (n > num_nodes):\n",
    "                continue\n",
    "            # Create two edges (forwards and backwards) with capacities based on\n",
    "            # histogram matching\n",
    "            hn = norm_hists[n]             # histogram for neighbor\n",
    "            g.add_edge(nodes[i], nodes[n], 20-cv2.compareHist(hi, hn, hist_comp_alg),\n",
    "                                           20-cv2.compareHist(hn, hi, hist_comp_alg))\n",
    "\n",
    "    # Match term: cost to FG/BG\n",
    "    for i,h in enumerate(norm_hists):\n",
    "        if i in fgbg_superpixels[0]:\n",
    "            g.add_tedge(nodes[i], 0, 1000) # FG - set high cost to BG\n",
    "        elif i in fgbg_superpixels[1]:\n",
    "            g.add_tedge(nodes[i], 1000, 0) # BG - set high cost to FG\n",
    "        else:\n",
    "            g.add_tedge(nodes[i], cv2.compareHist(fgbg_hists[0], h, hist_comp_alg),cv2.compareHist(fgbg_hists[1], h, hist_comp_alg))\n",
    "\n",
    "    g.maxflow()\n",
    "    return g.get_grid_segments(nodes)\n",
    "\n",
    "# Function to handle mouse movements\n",
    "def draw_segment(event,x,y,flags,param):\n",
    "    global refPt, haveBG, haveFG, newMask\n",
    "\n",
    "    # Store the point when mouse dragging is started\n",
    "    if event == cv2.EVENT_MBUTTONDOWN or event == cv2.EVENT_LBUTTONDOWN:\n",
    "        refPt = [(x,y)]\n",
    "    # Draw line between mouse start and end. Update segmentation if both foreground and background markings are present.\n",
    "    if event == cv2.EVENT_MBUTTONUP or event == cv2.EVENT_LBUTTONUP:\n",
    "        refPt.append((x, y))\n",
    "        if fgColor == True:    # Mode = foreground\n",
    "            haveFG = True\n",
    "            cv2.line(image,refPt[0],refPt[1],(0,0,255),5)    # Draw Red line\n",
    "            cv2.line(newMask,refPt[0],refPt[1],(0,0,255),5)\n",
    "        elif bgColor == True:    # Mode = background\n",
    "            haveBG = True\n",
    "            cv2.line(image,refPt[0],refPt[1],(255,0,0),5)    # Draw Blue line\n",
    "            cv2.line(newMask,refPt[0],refPt[1],(255,0,0),5)\n",
    "        # If both foreground and backgroudn markings are available then generate the binary image with foreground and background\n",
    "        if haveBG == True and haveFG == True:\n",
    "            #Calculating Foreground and Background superpixels using marking\n",
    "            fg_segments, bg_segments = find_superpixels_under_marking(newMask, superpixels)\n",
    "            \n",
    "            #Calculating color histograms for FG\n",
    "            fg_cumulative_hist = cumulative_histogram_for_superpixels(fg_segments, color_hists)\n",
    "            \n",
    "            #Calculating color histograms for BG\n",
    "            bg_cumulative_hist = cumulative_histogram_for_superpixels(bg_segments, color_hists)\n",
    "            \n",
    "            #Construct a graph that takes into account superpixel-to-superpixel interaction (smoothness term), as well as superpixel-FG/BG interaction (match term)\n",
    "            graph_cut = do_graph_cut([fg_cumulative_hist,bg_cumulative_hist], [fg_segments,bg_segments], norm_hists, neighbors)\n",
    "            \n",
    "            mask = pixels_for_segment_selection(superpixels, np.nonzero(graph_cut))\n",
    "            # mask is bool, conver 1 to 255 and 0 will remain 0 for displaying purpose\n",
    "            mask = np.uint8(mask * 255)\n",
    "            output_name = Path(input_path).stem\n",
    "            #thr, bin = cv2.threshold(mask, 120, 255.0, cv2.THRESH_BINARY);\n",
    "            #print(\"threshold: \",thr)\n",
    "            mask = cv2.resize(mask, (img.shape[1], img.shape[0]))\n",
    "            result = cv2.bitwise_or(img, img, mask=mask);\n",
    "            res_name = output_name+'_segment.png'\n",
    "            path = '/home/captainvish/Project/Final Year/Output'\n",
    "            cv2.imwrite(os.path.join(path,res_name), result)\n",
    "            # Show binary image\n",
    "            cv2.imshow(res_name,mask)\n",
    "            pathmask = '/home/captainvish/Project/Final Year/Mask'\n",
    "            mask_name = output_name+'_mask.png'\n",
    "            cv2.imwrite(os.path.join(pathmask,mask_name), mask)\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    \n",
    "    list_dr = '/home/captainvish/Project/Final Year/Input'\n",
    "    for image in os.listdir(list_dr):  \n",
    "        # Reading image\n",
    "        input_path = os.path.join(list_dr, image)\n",
    "        img = cv2.imread(input_path, cv2.IMREAD_COLOR)\n",
    "        imgname = image\n",
    "        height, width,band = img.shape\n",
    "        print(\"Height:  \" + str(height))\n",
    "        print(\"Width:   \" + str(width))\n",
    "        \n",
    "        #Calculating SLIC over image\n",
    "        centers, color_hists, superpixels, neighbors = superpixels_histograms_neighbors(img)\n",
    "        norm_hists = normalize_histograms(color_hists)\n",
    "\n",
    "        image = img.copy()\n",
    "\n",
    "        cv2.namedWindow(imgname)\n",
    "        cv2.setMouseCallback(imgname,draw_segment)\n",
    "        fgColor = True\n",
    "        bgColor = False\n",
    "        haveFG = False\n",
    "        haveBG = False\n",
    "        #Create a mask with only markings used in segmentation in original image.\n",
    "        newMask = np.zeros((img.shape), np.uint8)\n",
    "        # Fill image with white color(set each pixel to red).\n",
    "        newMask[:] = (255, 255, 255)\n",
    "        while(1):\n",
    "            cv2.imshow(imgname,image)\n",
    "            key = cv2.waitKey(1) & 0xFF\n",
    "            # Stop is user press 'q'\n",
    "            if key == ord(\"q\"):\n",
    "                break;\n",
    "            if key == ord(\"s\"):\n",
    "                mask_name = Path(input_path).stem + '_mark.png'\n",
    "                path = '/home/captainvish/Project/Final Year/Marking'\n",
    "                cv2.imwrite(os.path.join(path,mask_name), newMask)\n",
    "                break;\n",
    "            # Switch to background mode if user press 'c'\n",
    "            if key == ord(\"b\"):\n",
    "                bgColor = True\n",
    "                fgColor = False\n",
    "            # Switch to foreground mode if user press 'c'\n",
    "            if key == ord(\"f\"):\n",
    "                fgColor = True\n",
    "                bgColor = False\n",
    "            # Reset markings if user press 'r'\n",
    "            if key == ord(\"r\"):\n",
    "                image = img.copy()\n",
    "                bgColor = False\n",
    "                fgColor = True\n",
    "                haveFG = False\n",
    "                haveBG = False\n",
    "                newMask[:] = (255, 255, 255)\n",
    "                cv2.destroyWindow(imgname)\n",
    "        cv2.destroyAllWindows()\n",
    "        \n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7a7f580",
   "metadata": {},
   "source": [
    "### 2. Yield Count using Watershed algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ab8dd810",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0 0 0 ... 0 0 0]\n",
      " [0 0 0 ... 0 0 0]\n",
      " [0 0 0 ... 0 0 0]\n",
      " ...\n",
      " [0 0 0 ... 0 0 0]\n",
      " [0 0 0 ... 0 0 0]\n",
      " [0 0 0 ... 0 0 0]]\n",
      "40 Arecanuts found in the image segment.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3991/811683421.py:22: FutureWarning: indices argument is deprecated and will be removed in version 0.20. To avoid this warning, please do not use the indices argument. Please see peak_local_max documentation for more details.\n",
      "  localMax = peak_local_max(D, indices=False, min_distance=27,labels=thresh)\n",
      "/home/captainvish/.local/lib/python3.8/site-packages/skimage/morphology/_deprecated.py:5: skimage_deprecation: Function ``watershed`` is deprecated and will be removed in version 0.19. Use ``skimage.segmentation.watershed`` instead.\n",
      "  def watershed(image, markers=None, connectivity=1, offset=None, mask=None,\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[  0   0   0 ...   0   0   0]\n",
      " [  0   0   0 ...   0   0   0]\n",
      " [  0   0   0 ...   0   0   0]\n",
      " ...\n",
      " [255   0   0 ...   0   0   0]\n",
      " [  0   0   0 ...   0   0   0]\n",
      " [  0   0   0 ...   0   0   0]]\n",
      "35 Arecanuts found in the image segment.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3991/811683421.py:22: FutureWarning: indices argument is deprecated and will be removed in version 0.20. To avoid this warning, please do not use the indices argument. Please see peak_local_max documentation for more details.\n",
      "  localMax = peak_local_max(D, indices=False, min_distance=27,labels=thresh)\n",
      "/home/captainvish/.local/lib/python3.8/site-packages/skimage/morphology/_deprecated.py:5: skimage_deprecation: Function ``watershed`` is deprecated and will be removed in version 0.19. Use ``skimage.segmentation.watershed`` instead.\n",
      "  def watershed(image, markers=None, connectivity=1, offset=None, mask=None,\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[255 255 255 ...   0   0   0]\n",
      " [255 255 255 ...   0   0   0]\n",
      " [255 255 255 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...   0   0   0]\n",
      " [  0   0   0 ...   0   0   0]\n",
      " [  0   0   0 ...   0   0   0]]\n",
      "34 Arecanuts found in the image segment.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3991/811683421.py:22: FutureWarning: indices argument is deprecated and will be removed in version 0.20. To avoid this warning, please do not use the indices argument. Please see peak_local_max documentation for more details.\n",
      "  localMax = peak_local_max(D, indices=False, min_distance=27,labels=thresh)\n",
      "/home/captainvish/.local/lib/python3.8/site-packages/skimage/morphology/_deprecated.py:5: skimage_deprecation: Function ``watershed`` is deprecated and will be removed in version 0.19. Use ``skimage.segmentation.watershed`` instead.\n",
      "  def watershed(image, markers=None, connectivity=1, offset=None, mask=None,\n"
     ]
    }
   ],
   "source": [
    "# import the necessary packages\n",
    "from skimage.feature import peak_local_max\n",
    "from skimage.morphology import watershed\n",
    "from scipy import ndimage\n",
    "import numpy as np\n",
    "import argparse\n",
    "import imutils\n",
    "import cv2\n",
    "import os\n",
    "\n",
    "path = '/home/captainvish/Project/Final Year/Output'\n",
    "for filename in os.listdir(path):\n",
    "    input_path = os.path.join(path, filename)\n",
    "    img = cv2.imread(input_path)\n",
    "    shifted = cv2.pyrMeanShiftFiltering(img, 21, 51)\n",
    "\n",
    "    # Otsu's thresholding\n",
    "    gray = cv2.cvtColor(shifted, cv2.COLOR_BGR2GRAY)\n",
    "    thresh = cv2.threshold(gray, 0, 255,cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]\n",
    "    print(thresh)\n",
    "    D = ndimage.distance_transform_edt(thresh)\n",
    "    localMax = peak_local_max(D, indices=False, min_distance=27,labels=thresh)\n",
    "    \n",
    "    # perform a connected component analysis on the local peaks,\n",
    "    # using 8-connectivity, then appy the Watershed algorithm\n",
    "    markers = ndimage.label(localMax, structure=np.ones((3, 3)))[0]\n",
    "    labels = watershed(-D, markers, mask=thresh)\n",
    "    print(\"{} Arecanuts found in the image segment.\".format(len(np.unique(labels)) - 1))\n",
    "    for label in np.unique(labels):\n",
    "        # if the label is zero, we are examining the 'background'\n",
    "        # so simply ignore it\n",
    "        if label == 0:\n",
    "            continue\n",
    "        # otherwise, allocate memory for the label region and draw\n",
    "        # it on the mask\n",
    "        mask = np.zeros(gray.shape, dtype=\"uint8\")\n",
    "        mask[labels == label] = 255\n",
    "        # detect contours in the mask and grab the largest one\n",
    "        cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)\n",
    "        cnts = imutils.grab_contours(cnts)\n",
    "        c = max(cnts, key=cv2.contourArea)\n",
    "        # draw a circle enclosing the object\n",
    "        ((x, y), r) = cv2.minEnclosingCircle(c)\n",
    "        #cv2.circle(img, (int(x), int(y)), int(r), (0, 255, 0), 2)\n",
    "        cv2.putText(img, \"#{}\".format(label), (int(x) - 10, int(y)),cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)\n",
    "    # show the output image\n",
    "    yield_name = Path(input_path).stem + '_count.png'\n",
    "    path1 = '/home/captainvish/Project/Final Year/Yield Count'\n",
    "    cv2.imwrite(os.path.join(path1,yield_name), img)\n",
    "    cv2.imshow(yield_name, img)\n",
    "    cv2.waitKey(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66d6cade",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}