themblem/detection/roi_fix.py
2024-09-01 21:51:50 +01:00

145 lines
5.6 KiB
Python

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File : roi_fix.py
@Contact : zpyovo@hotmail.com
@License : (C)Copyright 2018-2019, Lab501-TransferLearning-SCUT
@Description :
@Modify Time @Author @Version @Desciption
------------ ------- -------- -----------
2024/2/27 00:02 Pengyu Zhang 1.0 None
'''
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
'''批量固定尺寸'''
import os
import cv2
from tqdm import tqdm
import numpy as np
def merge_overlapping_rectangles(rectangles):
if len(rectangles) <= 1:
return rectangles
merged_rectangles = []
for rect in sorted(rectangles, key=lambda x: x[2]*x[3], reverse=True):
x, y, w, h = rect
if not any((x < x2 + w2 and x + w > x2 and y < y2 + h2 and y + h > y2) for x2, y2, w2, h2 in merged_rectangles):
merged_rectangles.append(rect)
return merged_rectangles
def roi_Inner_detector(input,):
# 定义结构元素(内核)
kernel = np.ones((3, 3), np.uint8)
# 开运算
opening = cv2.morphologyEx(input, cv2.MORPH_OPEN, kernel)
# 定义结构元素(内核)
kernel = np.ones((3, 3), np.uint8)
# 闭运算
closing = cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel)
# 使用 Canny 边缘检测
edges = cv2.Canny(closing, 50, 255, apertureSize=3)
# 寻找轮廓
contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
# 临时存储矩形
temp_rectangles = []
# 识别矩形轮廓
for contour in contours:
# 轮廓近似
epsilon = 0.001 * cv2.arcLength(contour, True)
approx = cv2.approxPolyDP(contour, epsilon, True)
if len(approx) >= 4:
x, y, w, h = cv2.boundingRect(approx)
aspect_ratio = w / float(h)
area_contour = cv2.contourArea(contour)
min_area_threshold = 0
# print('aspect_ratio',aspect_ratio)
# print('area_contour',area_contour)
if aspect_ratio >= 0.80 and aspect_ratio < 1.2 and 6500 < area_contour < 13000:
area_bounding_rect = w * h
area_ratio = area_contour / area_bounding_rect
# print('aspect_ratio', area_ratio)
# print('area_ratio', area_ratio)
# print('area_contour', area_contour)
if 0.80 < area_ratio < 1.05:
temp_rectangles.append((x, y, w, h))
# 合并重叠的方框
merged_rectangles = merge_overlapping_rectangles(temp_rectangles)
return merged_rectangles
input_dir = '/project/dataset/QR2024/org100/orange-0-roi-fix'
output_dir = '/project/dataset/QR2024/org100/orange-0-roi-fix2fix'
size = 128
# 创建输出目录
os.makedirs(output_dir, exist_ok=True)
# 遍历输入目录中的图片
for filename in tqdm(os.listdir(input_dir), desc='Processing'):
if filename.endswith('.jpg') or filename.endswith('.png'):
input_path = os.path.join(input_dir, filename)
image = cv2.imread(input_path)
height, width, _ = image.shape
if height < size:
image = cv2.resize(np.array(image), (size, size), interpolation=cv2.INTER_AREA)
else:
image = cv2.resize(np.array(image), (size, size), interpolation=cv2.INTER_LINEAR)
image_gay = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
_, binary_image_A = cv2.threshold(image_gay, 50, 255, cv2.THRESH_BINARY)
merged_rectangles = roi_Inner_detector(binary_image_A)
if len(merged_rectangles) == 0 :
_, binary_image_A = cv2.threshold(image_gay, 150, 255, cv2.THRESH_BINARY)
merged_rectangles = roi_Inner_detector(binary_image_A)
# print('150')
if len(merged_rectangles) == 0 :
_, binary_image_A = cv2.threshold(image_gay, 100, 255, cv2.THRESH_BINARY)
merged_rectangles = roi_Inner_detector(binary_image_A)
# print('100')
if len(merged_rectangles) == 0 :
_, binary_image_A = cv2.threshold(image_gay, 125, 255, cv2.THRESH_BINARY)
merged_rectangles = roi_Inner_detector(binary_image_A)
# print('120')
if len(merged_rectangles) == 0 :
_, binary_image_A = cv2.threshold(image_gay, 75, 255, cv2.THRESH_BINARY)
merged_rectangles = roi_Inner_detector(binary_image_A)
if len(merged_rectangles) == 0:
_, binary_image_A = cv2.threshold(image_gay, 85, 255, cv2.THRESH_BINARY)
merged_rectangles = roi_Inner_detector(binary_image_A)
if len(merged_rectangles) == 0:
_, binary_image_A = cv2.threshold(image_gay, 115, 255, cv2.THRESH_BINARY)
merged_rectangles = roi_Inner_detector(binary_image_A)
# print('75')
# if len(merged_rectangles) == 0 :
# _, binary_image_A = cv2.threshold(image_gay, 0, 255, cv2.THRESH_BINARY)
# merged_rectangles = roi_Inner_detector(binary_image_A)
if merged_rectangles:
outermost_rectangle = max(merged_rectangles, key=lambda x: x[2] * x[3])
x, y, w, h = outermost_rectangle
image_crope = image[y + 1:y - 1 + h, x + 1:x - 1 + w]
image_crope = cv2.resize(image_crope, (32, 32), interpolation=cv2.INTER_AREA)
output_path = os.path.join(output_dir, filename)
cv2.imwrite(output_path, image_crope)
else:
print("failed",filename)
# output_path = os.path.join(output_dir, filename)
# cv2.imwrite(output_path, image)
# 保存裁剪后的图像
print("裁剪完成!")