count_eggs_sahi.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Count eggs using roboflow trained roboflow 3.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import supervision as sv\n",
    "from inference import get_roboflow_model\n",
    "import torch\n",
    "import onnxruntime as rt\n",
    "import cv2\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "import multiprocessing as mp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def callback(image_slice: np.ndarray) -> sv.Detections:\n",
    "    results = model.infer(image_slice, confidence=0.45, iou_threshold=0.5)[0] # Reset confidence level here\n",
    "    return sv.Detections.from_inference(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_eggs(image_path):\n",
    "    # Read image path\n",
    "    image = cv2.imread(image_path)\n",
    "\n",
    "    # Decode the QR code in the image\n",
    "    qcd = cv2.QRCodeDetector()\n",
    "\n",
    "    retval, decoded_info, points, straight_qrcode = qcd.detectAndDecodeMulti(image)\n",
    "    points\n",
    "\n",
    "    if not retval:\n",
    "        decoded_info = [os.path.basename(image_path)]\n",
    "    else:\n",
    "        # Calculate the center of the square\n",
    "        center = np.mean(points, axis=1)\n",
    "\n",
    "        # Expand the points by the specified factors\n",
    "        scale_x = 1.6\n",
    "        scale_y = 2.25\n",
    "        expanded_points = np.copy(points)\n",
    "        expanded_points[:, :, 0] = scale_x * (points[:, :, 0] - center[:, 0]) + center[:, 0]\n",
    "        expanded_points[:, :, 1] = scale_y * (points[:, :, 1] - center[:, 1]) + center[:, 1]\n",
    "\n",
    "        # Convert points to integer coordinates\n",
    "        expanded_points = expanded_points.astype(int)\n",
    "\n",
    "        # Create a mask with the same dimensions as the image\n",
    "        mask = np.ones(image.shape[:2], dtype=np.uint8) * 255\n",
    "\n",
    "        # Fill the mask with the expanded points\n",
    "        cv2.fillPoly(mask, expanded_points, 0)\n",
    "\n",
    "        # Apply the mask to the image to remove the specified points\n",
    "        image = cv2.bitwise_and(image, image, mask=mask)\n",
    "    \n",
    "    # Slice Image and run inference using callback function\n",
    "    slicer = sv.InferenceSlicer(callback=callback, slice_wh=(640, 640))\n",
    "    sliced_detections = slicer(image=image)\n",
    "\n",
    "    # Save output image\n",
    "    box_annotator = sv.BoxAnnotator()\n",
    "    # You can also use sv.MaskAnnotator() for instance segmentation models\n",
    "    # mask_annotator = sv.MaskAnnotator()\n",
    "    annotated_image = box_annotator.annotate(\n",
    "        scene=image.copy(), detections=sliced_detections)\n",
    "    # Convert the annotated image from numpy array to PIL Image\n",
    "    image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)\n",
    "    annotated_image_pil = Image.fromarray(image_rgb)\n",
    "    annotated_image_resized = annotated_image_pil.resize((4000, 6000))\n",
    "    annotated_image_resized.save(f'User/path/to/images/{decoded_info[0]}.png')\n",
    "\n",
    "    # Save information to create a pandas df for export\n",
    "    class_names = sliced_detections['class_name']\n",
    "    class_id = sliced_detections.class_id\n",
    "    confidence = sliced_detections.confidence\n",
    "\n",
    "    # save objects into a single df\n",
    "    df = (pd.DataFrame({'ID': class_names, 'class_id': class_id, 'confidence': confidence})\n",
    "        .value_counts(subset = 'ID')\n",
    "        .to_frame(decoded_info[0])\n",
    "        .T.assign(total_eggs = len(sliced_detections)))\n",
    "    \n",
    "    return df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load images\n",
    "directory = \"User/path/to/images\"\n",
    "image_paths = [os.path.join(directory, filename) for filename in os.listdir(directory) if filename.endswith('.JPG')]\n",
    "\n",
    "len(image_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_roboflow_model(model_id=\"egg_training-bi/1\", api_key=\"l6XPyOniqM4Ecq129cpf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For loop run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Multi Image run for model\n",
    "results = []\n",
    "\n",
    "# Process images sequentially using a for loop\n",
    "for image_path in image_paths:\n",
    "    try:\n",
    "        result = count_eggs(image_path)\n",
    "        results.append(result)\n",
    "        print(f\"Processed {image_path}\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing {image_path}: {e}\")\n",
    "        continue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multithread run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_image(image_path):\n",
    "    try:\n",
    "        result = count_eggs(image_path)\n",
    "        return result\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing {image_path}: {e}\")\n",
    "        return None\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # List of image paths\n",
    "    directory = \"/Users/aja294/Library/CloudStorage/Box-Box/Breeding Insight/Species/Salmonids/Idaho Group/Trout/Image Analysis/final_images_2024/eggs\"\n",
    "    image_paths = [os.path.join(directory, filename) for filename in os.listdir(directory) if filename.endswith('.JPG')]\n",
    "\n",
    "    # Number of worker processes\n",
    "    num_workers = mp.cpu_count() // 2  # Use half of the available CPU cores\n",
    "\n",
    "    # Create a pool of worker processes\n",
    "    with mp.Pool(num_workers) as pool:\n",
    "        # Map the image paths to the process_image function\n",
    "        results = pool.map(process_image, image_paths)\n",
    "\n",
    "    # Filter out None results (in case of errors)\n",
    "    results = [result for result in results if result is not None]\n",
    "\n",
    "    # Print the results\n",
    "    print(\"Results:\", results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combine results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine all DataFrames into a single DataFrame\n",
    "final_df = pd.concat(results)\n",
    "\n",
    "# Print the final DataFrame\n",
    "print(final_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_df.to_csv('result/path', index=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "egg_counter-env",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}