aug.py for generating miising screws and bad screws

model/aug.ipynb

@@ -0,0 +1,65 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import cv2\n",
+    "import numpy as np\n",
+    "import random"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(0.9029411764705882, 0.7752941176470588, 0.023529411764705882, 0.023529411764705882)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Загрузка изображений ноутбука, дефекта и маски\n",
+    "from aug import apply_defect\n",
+    "\n",
+    "laptop_img = cv2.imread('res/laptop.jpg')\n",
+    "defect_img = cv2.imread('res/defect2.png')  # Изображение дефекта (например, шуруп)\n",
+    "mask_img = cv2.imread('res/mask.jpg', 0)  # Маска (ч/б изображение), указывающая области для дефектов\n",
+    "\n",
+    "# Применение дефекта\n",
+    "result, bbox = apply_defect(laptop_img, defect_img, mask_img)\n",
+    "\n",
+    "# Сохранение результата\n",
+    "cv2.imwrite('res/laptop_with_defect.jpg', result)\n",
+    "print(bbox)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "torch-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.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

model/aug.py

@@ -0,0 +1,102 @@
+import cv2
+import numpy as np
+import random
+
+# Функция для регулировки яркости и контрастности
+def adjust_brightness_contrast(image, brightness=0, contrast=0):
+    img = cv2.convertScaleAbs(image, alpha=contrast, beta=brightness)
+    return img
+
+# Функция для наложения дефекта на изображение ноутбука и возврата bbox
+def apply_defect(laptop_image, defect_image, mask, radius=20, brightness=0, contrast=1.0):
+    h, w = laptop_image.shape[:2]
+    
+    # Генерация случайной позиции внутри маски
+    while True:
+        x = random.randint(radius, w - radius)
+        y = random.randint(radius, h - radius)
+        if mask[y, x] > 0:
+            break
+    
+    # Создание новой маски в форме круга с размытыми краями
+    new_mask = np.zeros((h, w), dtype=np.uint8)
+    cv2.circle(new_mask, (x, y), radius, 255, -1)
+    # new_mask = cv2.GaussianBlur(new_mask, (51, 51), 20)  # Размытие маски
+
+    # Масштабирование дефекта для улучшенной видимости
+    defect_resized = cv2.resize(defect_image, (2 * radius, 2 * radius))
+
+    # Случайный угол поворота от 0 до 360 градусов
+    angle = random.uniform(0, 360)
+    
+    # Центр для поворота
+    center = (radius, radius)
+    
+    # Матрица поворота
+    M = cv2.getRotationMatrix2D(center, angle, 1.0)
+    
+    # Применение поворота
+    defect_rotated = cv2.warpAffine(defect_resized, M, (2 * radius, 2 * radius), flags=cv2.INTER_LINEAR)
+
+    # Создание маски для черных пикселей
+    black_mask = np.all(defect_rotated == [0, 0, 0], axis=-1).astype(np.uint8)
+
+    kernel = np.ones((2, 2), np.uint8)  # Ядро 3x3, чтобы расширить на 2 пикселя (делаем 2 итерации)
+    black_mask_dilated = cv2.dilate(black_mask, kernel, iterations=2)
+
+    # Вычитание черных областей из общей маски
+    mask_region = new_mask[y - radius:y + radius, x - radius:x + radius]
+    mask_region[black_mask_dilated == 1] = 0
+
+    
+
+    # Регулировка яркости и контрастности
+    defect_rotated = adjust_brightness_contrast(defect_rotated, brightness, contrast)
+    
+    # Координаты области для вставки дефекта
+    x1, y1 = x - radius, y - radius
+    x2, y2 = x + radius, y + radius
+    
+    # Создание альфа-канала для плавного наложения дефекта
+    alpha = mask_region / 255.0
+
+    # Наложение дефекта с учетом маски (игнорируем черные области)
+    for c in range(0, 3):
+        laptop_image[y1:y2, x1:x2, c] = (alpha * defect_rotated[:, :, c] +
+                                         (1 - alpha) * laptop_image[y1:y2, x1:x2, c])
+
+
+     # Расширение маски на 1 пиксель
+    expanded_mask = cv2.dilate(new_mask, np.ones((2, 2), np.uint8), iterations=1)
+
+    # Сжатие маски на 2 пикселя
+    eroded_mask = cv2.erode(new_mask, np.ones((3, 3), np.uint8), iterations=1)
+
+    # Вычитание сжатой маски из расширенной маски
+    blurred_mask = cv2.subtract(expanded_mask, eroded_mask)
+
+    # Размытие laptop_image по краевой маске
+    blurred_laptop_image = cv2.GaussianBlur(laptop_image, (21, 21), 10)  # Размытие всего изображения
+    laptop_image = np.where(blurred_mask[:, :, None] > 0, blurred_laptop_image, laptop_image)  # Заменяем область размытым изображением
+
+
+    # Размытие laptop_image по краевой маске
+    blurred_laptop_image = cv2.GaussianBlur(laptop_image, (21, 21), 10)  # Размытие всего изображения
+    laptop_image = np.where(blurred_mask[:, :, None] > 0, blurred_laptop_image, laptop_image)  # Заменяем область размытым изображением
+
+    # Вычисление bbox в формате YOLO
+    bbox_x_center = (x1 + x2) / 2 / w  # нормализованный x центра
+    bbox_y_center = (y1 + y2) / 2 / h  # нормализованный y центра
+    bbox_width = (x2 - x1) / w         # нормализованная ширина
+    bbox_height = (y2 - y1) / h        # нормализованная высота
+
+    return laptop_image, (bbox_x_center, bbox_y_center, bbox_width, bbox_height)
+
+def mask_magic(mask, ex = (3,3), er = (3,3)):
+    # Расширение маски на 1 пиксель
+    expanded_mask = cv2.dilate(mask, np.ones(ex, np.uint8), iterations=1)
+
+    # Сжатие маски на 2 пиксел��
+    eroded_mask = cv2.erode(mask, np.ones(er, np.uint8), iterations=1)
+
+    return cv2.subtract(expanded_mask, eroded_mask)

model/yolo11n.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ebbc80d4a7680d14987a577cd21342b65ecfd94632bd9a8da63ae6417644ee1
+size 5613764