{ "cells": [ { "cell_type": "code", "execution_count": 10, "id": "6dac49da-4935-4f90-b262-2db73f4fbc77", "metadata": {}, "outputs": [], "source": [ "import os\n", "import numpy as np\n", "import cv2\n", "import matplotlib.pyplot as plt\n", "\n", "def read_txt(file_path, image):\n", " with open(file_path, 'r') as f:\n", " lines = f.readlines()\n", " image_height, image_width = image.shape[:2]\n", " edge = []\n", " node = []\n", " for line in lines:\n", " line = line.strip().split(' ')\n", " category = int(line[0])\n", " x_norm, y_norm, w_norm, h_norm = map(float,line[1:5])\n", " \n", " x = x_norm * image_width\n", " y = y_norm * image_height\n", " w = w_norm * image_width\n", " h = h_norm * image_height\n", " \n", "\n", " if category == 0:\n", " if w < h:\n", " edge1 = (x, y - h/2)\n", " edge2 = (x, y + h/2)\n", " edge.append([(x,y,w,h), edge1, edge2])\n", " elif w>= h:\n", " edge1 = (x - w/2, y)\n", " edge2 = (x + w/2, y)\n", " edge.append([(x,y,w,h), edge1, edge2])\n", " elif category in [1, 2, 3, 5]:\n", " t = (x, y + h/2)\n", " b = (x, y - h/2)\n", " l = (x - w/2, y)\n", " r = (x + w/2, y)\n", " node.append([(x,y,w,h), t, b, l, r])\n", " \n", " return edge, node\n", "\n", "def calculate_distance(point1, point2):\n", " x1, y1 = point1\n", " x2, y2 = point2\n", " return np.sqrt((x1 - x2)**2 + (y1 - y2)**2)\n", "\n", "def find_closest_node(edge, node):\n", " results = []\n", " for edge_box in edge:\n", " edge1, edge2 = edge_box[1], edge_box[2]\n", " # print('==============',edge1, edge2)\n", " min_distance1 = float('inf')\n", " min_distance2 = float('inf')\n", " \n", " closest_node = None\n", " for node_box in node:\n", " # print(node_box)\n", " for i in range(4):\n", " node_point = node_box[1 + i]\n", "\n", "\n", " distance1 = calculate_distance(edge1, node_point)\n", " distance2 = calculate_distance(edge2, node_point)\n", "\n", " \n", " if distance1 < min_distance1:\n", " min_distance1 = distance1\n", " closest_node1 = node_box\n", " \n", " if distance2 < min_distance2:\n", " min_distance2 = distance2\n", " closest_node2 = node_box\n", " # print('=======================',closest_node)\n", " results.append(('edge_box',edge_box, 'closest_node1',closest_node1, 'closest_node2',closest_node2))\n", " return results\n", "\n", "def xywh_to_xyxy(x, y, w, h):\n", " xmin = x - w/2\n", " ymin = y - h/2\n", " xmax = x + w/2\n", " ymax = y + h/2\n", " return int(xmin), int(ymin), int(xmax), int(ymax)\n", "\n", "def crop_xyxy(boxes_0,boxes_1,boxes_2):\n", " x0,y0,w0,h0 = boxes_0 \n", " x1,y1,w1,h1 = boxes_1\n", " x2,y2,w2,h2 = boxes_2\n", " xmin0,ymin0,xmax0,ymax0 = xywh_to_xyxy(x0,y0,w0,h0)\n", " xmin1,ymin1,xmax1,ymax1 = xywh_to_xyxy(x1,y1,w1,h1)\n", " xmin2,ymin2,xmax2,ymax2 = xywh_to_xyxy(x2,y2,w2,h2)\n", " crop_xmin = min(xmin0,xmin1,xmin2)\n", " crop_ymin = min(ymin0,ymin1,ymin2)\n", " crop_xmax = max(xmax0,xmax1,xmax2)\n", " crop_ymax = max(ymax0,ymax1,ymax2)\n", " return int(crop_xmin), int(crop_ymin), int(crop_xmax), int(crop_ymax)\n", "\n", "\n", "if __name__ == '__main__':\n", " file_dir = \"D:/jupyter/block_diagram/yolov5/runs/detect/exp9/labels/\"\n", " output_dir = \"D:/jupyter/block_diagram/output_images/\"\n", " image_dir = \"D:/jupyter/block_diagram/input_img/\"\n", " if not os.path.exists(output_dir):\n", " os.makedirs(output_dir)\n", "\n", " for txt_file in os.listdir(file_dir):\n", " if txt_file !='.ipynb_checkpoints': \n", " file_path = file_dir + txt_file\n", " img_name = txt_file[:-4]\n", " # print(img_name)\n", " image_path = image_dir + img_name + '.PNG'\n", "\n", " img = cv2.imread(image_path)\n", " edge, node = read_txt(file_path, img)\n", " result = find_closest_node(edge, node)\n", " for i in range(len(result)):\n", " boxes = [result[i][1][0],result[i][3][0],result[i][5][0]]\n", " crop_xmin, crop_ymin, crop_xmax, crop_ymax = crop_xyxy(boxes[0],boxes[1],boxes[2])\n", " # print(crop_xmin, crop_ymin, crop_xmax, crop_ymax)\n", " # print('===============')\n", " # 创建白色图像\n", " new_img = 255*np.ones(img.shape, img.dtype)\n", "\n", " # 循环拷贝三个边界框到新图像上\n", " for j, box in enumerate(boxes):\n", " x,y,w,h = box\n", " w = w+10\n", " h = h+10\n", " \n", "# box = xywh_to_xyxy(x, y, w, h)\n", "# box_xmin = []\n", "# box_ymin = []\n", "# box_xmax = []\n", "# box_ymax = []\n", "# box_xmin.append(box[0])\n", "# box_ymin.append(box[1])\n", "# box_xmax.append(box[2])\n", "# box_ymax.append(box[3])\n", " \n", " \n", "# roi = img[box[1]:box[3], box[0]:box[2]]\n", "# new_img[box[1]:box[3], box[0]:box[2]] = roi\n", " \n", "# box_combined = [min(box_xmin), min(box_ymin), max(box_xmax), max(box_ymax)]\n", "\n", "# img_cropped = new_img[box_combined[1]:box_combined[3], box_combined[0]:box_combined[2]]\n", "\n", " xmin, ymin, xmax, ymax = xywh_to_xyxy(x, y, w, h)\n", " \n", " roi = img[ymin:ymax, xmin:xmax]\n", " new_img[ymin:ymax, xmin:xmax] = roi\n", "\n", " \n", " croped = new_img[crop_ymin:crop_ymax, crop_xmin:crop_xmax]\n", "\n", " # 保存新图像\n", " # cv2.imwrite(output_dir+'{}_{:04d}.jpg'.format(img_name,i), new_img)\n", " cv2.imwrite(output_dir+'{}_{:04d}.png'.format(img_name,i), croped)\n" ] }, { "cell_type": "code", "execution_count": 8, "id": "c03cb46f-47b3-4dd6-b36d-6fdbe5dcbbe7", "metadata": {}, "outputs": [], "source": [ "import os\n", "import shutil\n", "\n", "# 定义要读取的文件夹路径\n", "folder_path = \"D:/jupyter/block_diagram/output_images\"\n", "\n", "# 遍历文件夹中的所有文件\n", "for file_name in os.listdir(folder_path):\n", " # 只处理图像文件\n", " if file_name.endswith(\".png\") or file_name.endswith(\".jpg\"):\n", " # 以 _ 分割文件名\n", " parts = file_name.split(\"_\")\n", " # 如果文件名不符合要求，则跳过\n", " if len(parts) < 2:\n", " continue\n", " # 提取前缀\n", " prefix = parts[0]\n", " # 检查文件名前缀的文件夹是否存在，如果不存在则创建\n", " if not os.path.exists(os.path.join(folder_path, prefix)):\n", " os.makedirs(os.path.join(folder_path, prefix))\n", " # 移动文件到对应的文件夹中\n", " shutil.move(\n", " os.path.join(folder_path, file_name),\n", " os.path.join(folder_path, prefix, file_name)\n", " )\n" ] }, { "cell_type": "code", "execution_count": null, "id": "bf047b84-52cb-44b7-a783-be8ff68aac9a", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "diadem", "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.16" } }, "nbformat": 4, "nbformat_minor": 5 }