diff --git a/DPL/dgd_class/model/best_model_11.14.19.30.onnx b/DPL/dgd_class/model/best_model_11.14.19.30.onnx deleted file mode 100644 index f1af532..0000000 Binary files a/DPL/dgd_class/model/best_model_11.14.19.30.onnx and /dev/null differ diff --git a/DPL/dgd_class/model/best_model_params_11.14.19.30.pt b/DPL/dgd_class/model/best_model_params_11.14.19.30.pt deleted file mode 100644 index b068a0a..0000000 Binary files a/DPL/dgd_class/model/best_model_params_11.14.19.30.pt and /dev/null differ diff --git a/DPL/main.py b/DPL/main.py new file mode 100644 index 0000000..e69e832 --- /dev/null +++ b/DPL/main.py @@ -0,0 +1,42 @@ +import argparse +import os +import sys + +import torch + +from model_cls.models import Model as ClsModel + +from pathlib import Path + + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[0] # YOLOv5 root directory +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative + +def main(): + # 命令行参数解析 + parser = argparse.ArgumentParser(description="Use an ONNX model for inference.") + + # 设置默认值，并允许用户通过命令行进行修改 + parser.add_argument('--input-dir', type=str, default='dataset/', help='Directory to input images') + parser.add_argument('--save-dir', type=str, default='detect', help='Directory to save output images') + parser.add_argument('--gpu', action='store_true', help='Use GPU for inference') + + args = parser.parse_args() + + # 设置设备 + device = torch.device('cuda' if args.gpu and torch.cuda.is_available() else 'cpu') + if args.gpu and not torch.cuda.is_available(): + print("GPU not available, switching to CPU.") + + # 加载模型 + model = ClsModel(model_path=args.weights, device=device) + + + + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/DPL/dgd_class/export.py b/DPL/model_cls/export.py similarity index 100% rename from DPL/dgd_class/export.py rename to DPL/model_cls/export.py diff --git a/DPL/dgd_class/models.py b/DPL/model_cls/models.py similarity index 100% rename from DPL/dgd_class/models.py rename to DPL/model_cls/models.py diff --git a/DPL/dgd_class/predict.py b/DPL/model_cls/predict.py similarity index 100% rename from DPL/dgd_class/predict.py rename to DPL/model_cls/predict.py diff --git a/DPL/model_cls/preprocess.py b/DPL/model_cls/preprocess.py new file mode 100644 index 0000000..bc616c3 --- /dev/null +++ b/DPL/model_cls/preprocess.py @@ -0,0 +1,196 @@ +import argparse +import os +import random +import shutil + +import cv2 +import numpy as np +import json +from shapely.geometry import Polygon, box +from shapely.affinity import translate + + +def create_dataset_from_folder(image_folder, label_folder, block_size, output_dir): + """ + 批量处理文件夹中的图片和对应标签，生成分类模型的数据集。 + + Args: + image_folder (str): 图片文件夹路径。 + label_folder (str): 标签文件夹路径，对应Labelme生成的JSON文件。 + block_size (tuple): 分块的尺寸 (width, height)。 + output_dir (str): 输出数据集的根目录。 + """ + # 创建输出文件夹 + has_label_dir = os.path.join(output_dir, "has_label") + no_label_dir = os.path.join(output_dir, "no_label") + os.makedirs(has_label_dir, exist_ok=True) + os.makedirs(no_label_dir, exist_ok=True) + + # 遍历图片文件夹 + for filename in os.listdir(image_folder): + if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tif')): + image_path = os.path.join(image_folder, filename) + label_path = os.path.join(label_folder, os.path.splitext(filename)[0] + ".json") + + # 检查标签文件是否存在 + if not os.path.exists(label_path): + print(f"Label file not found for image: {filename}") + continue + + print(f"Processing {filename}...") + process_single_image(image_path, label_path, block_size, has_label_dir, no_label_dir) + + +def process_single_image(image_path, label_path, block_size, has_label_dir, no_label_dir): + """ + 处理单张图片并分块保存到对应的文件夹。 + + Args: + image_path (str): 图片路径。 + label_path (str): 标签路径。 + block_size (tuple): 分块的尺寸 (width, height)。 + has_label_dir (str): 包含标注的分块保存目录。 + no_label_dir (str): 无标注的分块保存目录。 + """ + # 加载图片 + image = cv2.imread(image_path) + img_height, img_width, _ = image.shape + + # 加载Labelme JSON文件 + with open(label_path, 'r', encoding='utf-8') as f: + label_data = json.load(f) + + # 提取多边形标注 + polygons = [] + for shape in label_data['shapes']: + if shape['shape_type'] == 'polygon': + points = shape['points'] + polygons.append(Polygon(points)) + + if roi: + x_min, y_min, x_max, y_max = roi + x_min, y_min = max(0, x_min), max(0, y_min) + x_max, y_max = min(img_width, x_max), min(img_height, y_max) + image = image[y_min:y_max, x_min:x_max] + img_height, img_width = y_max - y_min, x_max - x_min + # 偏移标注的多边形 + polygons = [translate(poly.intersection(box(x_min, y_min, x_max, y_max)), -x_min, -y_min) for poly in polygons] + + # 分割图片并保存到对应的文件夹 + block_width, block_height = block_size + block_id = 0 + base_name = os.path.splitext(os.path.basename(image_path))[0] + for y in range(0, img_height, block_height): + for x in range(0, img_width, block_width): + # 当前分块的边界框 + block_polygon = box(x, y, x + block_width, y + block_height) + + # 判断是否与任何标注的多边形相交 + contains_label = any(poly.intersects(block_polygon) for poly in polygons) + + # 裁剪当前块 + block = image[y:y + block_height, x:x + block_width] + + # 保存到对应文件夹 + folder = has_label_dir if contains_label else no_label_dir + block_filename = os.path.join(folder, f"{base_name}_block_{block_id}.jpg") + cv2.imwrite(block_filename, block) + block_id += 1 + + print(f"Saved {block_filename} to {'has_label' if contains_label else 'no_label'} folder.") + + +def split_dataset(input_dir, output_dir, train_ratio=0.7, val_ratio=0.2, test_ratio=0.1): + """ + 将分类后的数据集划分为 train、val 和 test 数据集，保持来源目录结构。 + + Args: + input_dir (str): 分类结果根目录，包含多个子文件夹（标签文件夹）。 + output_dir (str): 输出数据集目录，将生成 train、val 和 test 文件夹。 + train_ratio (float): train 集比例或固定数量。 + val_ratio (float): val 集比例或固定数量。 + test_ratio (float): test 集比例或固定数量。 + """ + # 定义输出子目录 + train_dir = os.path.join(output_dir, "train") + val_dir = os.path.join(output_dir, "val") + test_dir = os.path.join(output_dir, "test") + + # 遍历所有子文件夹（标签文件夹） + for category in os.listdir(input_dir): + category_dir = os.path.join(input_dir, category) + if not os.path.isdir(category_dir): # 忽略非文件夹 + continue + + # 为当前类别在 train/val/test 创建相同的子文件夹结构 + os.makedirs(os.path.join(train_dir, category), exist_ok=True) + os.makedirs(os.path.join(val_dir, category), exist_ok=True) + os.makedirs(os.path.join(test_dir, category), exist_ok=True) + + # 获取当前类别下的所有文件 + files = os.listdir(category_dir) + random.shuffle(files) + + # 计算分割点 + total_files = len(files) + if train_ratio < 1: + train_count = int(total_files * train_ratio) + val_count = int(total_files * val_ratio) + test_count = total_files - train_count - val_count + else: + train_count = int(train_ratio) + val_count = int(val_ratio) + test_count = int(test_ratio) + + # 确保不超过文件总数 + train_count = min(train_count, total_files) + val_count = min(val_count, total_files - train_count) + test_count = min(test_count, total_files - train_count - val_count) + + print(f"Category {category}: {train_count} train, {val_count} val, {test_count} test files") + + # 划分数据集 + train_files = files[:train_count] + val_files = files[train_count:train_count + val_count] + test_files = files[train_count + val_count:] + + # 复制文件到对应的子目录 + for file in train_files: + shutil.copy(os.path.join(category_dir, file), os.path.join(train_dir, category, file)) + for file in val_files: + shutil.copy(os.path.join(category_dir, file), os.path.join(val_dir, category, file)) + for file in test_files: + shutil.copy(os.path.join(category_dir, file), os.path.join(test_dir, category, file)) + + print(f"Category {category} processed. Train: {len(train_files)}, Val: {len(val_files)}, Test: {len(test_files)}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Preprocess images to required shapes") + + # 设置默认值，并允许用户通过命令行进行修改 + parser.add_argument('--img-dir', type=str, default=r'..\dataset\dgd\test_img', help='Directory to input images') + parser.add_argument('--label-dir', type=str, default=r'..\dataset\dgd\test_img', help='Directory to input labels') + + parser.add_argument('--output-dir', type=str, default=r'..\dataset\dgd\runs', help='Directory to save output images') + parser.add_argument('--roi', type=int, nargs=4, default=None, help='ROI region (x_min y_min x_max y_max)') + parser.add_argument('--train-ratio', type=float, default=0.7, help='Train set ratio or count') + parser.add_argument('--val-ratio', type=float, default=0.2, help='Validation set ratio or count') + parser.add_argument('--test-ratio', type=float, default=0.1, help='Test set ratio or count') + + args = parser.parse_args() + + # 输入文件夹路径 + image_folder = args.img_dir + label_folder = args.label_dir + # 输出路径 + output_dir = args.output_dir + + # 分块大小 + block_size = (170, 170) # 替换为希望的分块尺寸 (宽, 高) + roi = tuple(args.roi) if args.roi else None + all_output = os.path.join(output_dir, 'all') + + # 批量生成数据集 + create_dataset_from_folder(image_folder, label_folder, block_size, all_output) + split_dataset(all_output, output_dir, args.train_ratio, args.val_ratio, args.test_ratio) \ No newline at end of file diff --git a/DPL/dgd_class/train.py b/DPL/model_cls/train.py similarity index 67% rename from DPL/dgd_class/train.py rename to DPL/model_cls/train.py index 4f0ce9c..d2601a9 100644 --- a/DPL/dgd_class/train.py +++ b/DPL/model_cls/train.py @@ -1,6 +1,8 @@ import os import argparse import time +from datetime import datetime + import torch import torch.nn as nn import torch.optim as optim @@ -10,7 +12,7 @@ import matplotlib.pyplot as plt # 设置 cudnn 优化 torch.backends.cudnn.benchmark = True -plt.ion() + def get_next_model_name(save_dir, base_name, ext): """ @@ -131,37 +133,58 @@ def train_model(model, criterion, optimizer, scheduler, dataloaders, dataset_siz return model -# 参数解析部分 -parser = argparse.ArgumentParser(description="Train a ResNet18 model on custom dataset.") -parser.add_argument('--data-dir', type=str, default='dataset', help='Path to the dataset directory.') -parser.add_argument('--model-save-dir', type=str, default='model', help='Directory to save the trained model.') -parser.add_argument('--batch-size', type=int, default=8, help='Batch size for training.') -parser.add_argument('--epochs', type=int, default=5, help='Number of training epochs.') -parser.add_argument('--lr', type=float, default=0.0005, help='Learning rate.') -parser.add_argument('--momentum', type=float, default=0.95, help='Momentum for SGD.') -parser.add_argument('--step-size', type=int, default=5, help='Step size for LR scheduler.') -parser.add_argument('--gamma', type=float, default=0.2, help='Gamma for LR scheduler.') -parser.add_argument('--num-classes', type=int, default=2, help='Number of output classes.') -parser.add_argument('--model-base-name', type=str, default='best_model', help='Base name for saved models.') - def main(args): """Main function to train the model based on command-line arguments.""" device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") dataloaders, dataset_sizes, class_names = get_data_loaders(args.data_dir, args.batch_size) - - model = models.resnet18(weights='IMAGENET1K_V1') - num_ftrs = model.fc.in_features - model.fc = nn.Linear(num_ftrs, args.num_classes) + if args.model_name == 'resnet18': + model = models.resnet18(weights='IMAGENET1K_V1') + num_ftrs = model.fc.in_features + model.fc = nn.Linear(num_ftrs, args.num_classes) + elif args.model_name == 'resnet50': + model = models.resnet50(weights='IMAGENET1K_V1') + num_ftrs = model.fc.in_features + model.fc = nn.Linear(num_ftrs, args.num_classes) + elif args.model_name == 'alexnet': + model = models.alexnet(pretrained=True) + num_ftrs = model.classifier[-1].in_features + model.classifier[-1] = nn.Linear(num_ftrs, args.num_classes) + elif args.model_name == 'vgg16': + model = models.vgg16(pretrained=True) + num_ftrs = model.classifier[-1].in_features + model.classifier[-1] = nn.Linear(num_ftrs, args.num_classes) + elif args.model_name == 'densenet': + model = models.densenet121(pretrained=True) + num_ftrs = model.classifier.in_features + model.classifier = nn.Linear(num_ftrs, args.num_classes) + elif args.model_name == 'inceptionv3': + model = models.inception_v3(pretrained=True) + num_ftrs = model.fc.in_features + model.fc = nn.Linear(num_ftrs, args.num_classes) model = model.to(device) + # training start time and model info + model_base_name = f"{args.model_name}_bs{args.batch_size}_ep{args.epochs}_{datetime.now().strftime('%y_%m_%d')}.pt" criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) scheduler = lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma) train_model(model, criterion, optimizer, scheduler, dataloaders, dataset_sizes, device, - args.epochs, args.model_save_dir, args.model_base_name) + args.epochs, args.model_save_dir, model_base_name) if __name__ == '__main__': + # 参数解析部分 + parser = argparse.ArgumentParser(description="Train a ResNet18 model on custom dataset.") + parser.add_argument('--model_name', type=str, default='resnet18', help='Which model to train as base model') + parser.add_argument('--data-dir', type=str, default='dataset', help='Path to the dataset directory.') + parser.add_argument('--model-save-dir', type=str, default='model', help='Directory to save the trained model.') + parser.add_argument('--batch-size', type=int, default=8, help='Batch size for training.') + parser.add_argument('--epochs', type=int, default=5, help='Number of training epochs.') + parser.add_argument('--lr', type=float, default=0.0005, help='Learning rate.') + parser.add_argument('--momentum', type=float, default=0.95, help='Momentum for SGD.') + parser.add_argument('--step-size', type=int, default=5, help='Step size for LR scheduler.') + parser.add_argument('--gamma', type=float, default=0.2, help='Gamma for LR scheduler.') + parser.add_argument('--num-classes', type=int, default=2, help='Number of output classes.') args = parser.parse_args() main(args) diff --git a/DPL/model_cls/utils.py b/DPL/model_cls/utils.py new file mode 100644 index 0000000..e69de29