main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ultralytics import YOLO\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class YOLOSegmentation:\n",
    "    def __init__(self, model_path):\n",
    "        self.model = YOLO(model_path)\n",
    "\n",
    "    def detect(self, img):\n",
    "        height, width, channels = img.shape\n",
    "\n",
    "        results = self.model.predict(source=img.copy(), save=False, save_txt=False)\n",
    "        result = results[0]\n",
    "        segmentation_contours_idx = []\n",
    "        for seg in result.masks.segments:\n",
    "            seg[:, 0] *= width\n",
    "            seg[:, 1] *= height\n",
    "            segment = np.array(seg, dtype=np.int32)\n",
    "            segmentation_contours_idx.append(segment)\n",
    "\n",
    "        bboxes = np.array(result.boxes.xyxy.cpu(), dtype=\"int\")\n",
    "        class_ids = np.array(result.boxes.cls.cpu(), dtype=\"int\")\n",
    "        scores = np.array(result.boxes.conf.cpu(), dtype=\"float\").round(2)\n",
    "        return bboxes, class_ids, segmentation_contours_idx, scores\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "\n",
    "\n",
    "img = cv2.imread(\"Path to image\")\n",
    "img = cv2.resize(img, None, fx=0.5, fy=0.5)\n",
    "\n",
    "\n",
    "ys = YOLOSegmentation(\"yolov8m-seg.pt\")\n",
    "\n",
    "bboxes, classes, segmentations, scores = ys.detect(img)\n",
    "\n",
    "#To draw contours, bbox and class of image uncomment it\n",
    "# for bbox, class_id, seg, score in zip(bboxes, classes, segmentations, scores):\n",
    "#     (x, y, x2, y2) = bbox\n",
    "#     # if class_id == 29:\n",
    "#     # cv2.rectangle(img, (x, y), (x2, y2), (255, 0, 0), 2)\n",
    "\n",
    "#     cv2.polylines(img, [seg], True, (0, 0, 255), 1)\n",
    "\n",
    "    # cv2.putText(\n",
    "    #     img, str(class_id), (x, y - 10),\n",
    "    #     cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 255), 2)\n",
    "\n",
    "# cv2.imshow(\"image\", img)\n",
    "# cv2.waitKey(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Все за пределами сегментированного изображения станет черным\n",
    "points = np.array(segmentations)\n",
    "\n",
    "mask = np.zeros(img.shape[:2], dtype=np.uint8)\n",
    "cv2.drawContours(mask, [points], -1, (255, 255, 255), -1, cv2.LINE_AA)\n",
    "cropped_img = cv2.bitwise_and(img, img, mask=mask)\n",
    "\n",
    "# Display cropped image if you want to\n",
    "# cv2.imshow(\"Cropped Image\", cropped_img)\n",
    "# cv2.waitKey(0)\n",
    "cv2.imwrite(('Where to store cropped image'), cropped_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image3 = cv2.imread(cropped_img) \n",
    "blurred_image = cv2.GaussianBlur(image3.copy(),(5,5),0)\n",
    "\n",
    "edges = cv2.Canny(blurred_image, 100, 160)\n",
    "\n",
    "plt.figure(figsize=[10,10])\n",
    "cv2.imwrite(('Where to store contours'), edges)\n",
    "plt.imshow(edges,cmap='Greys_r');plt.title(\"Edges Image\");plt.axis(\"off\")\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}