diff --git a/research/deploy.sh b/research/deploy.sh index a324639..11c4e7a 100755 --- a/research/deploy.sh +++ b/research/deploy.sh @@ -1,7 +1,7 @@ #!/bin/bash set -e -DOCKER_IMAGE=registry.gitlab.com/euphon/themblem:research-$(git rev-parse --short HEAD) +DOCKER_IMAGE=registry.cn-shenzhen.aliyuncs.com/emblem/themblem:research-$(git rev-parse --short HEAD) docker build --network=host -t $DOCKER_IMAGE . docker push $DOCKER_IMAGE diff --git a/research/models/resnet18_20250325_114510_94.56.pth b/research/models/resnet18_20250325_114510_94.56.pth new file mode 100644 index 0000000..5835930 Binary files /dev/null and b/research/models/resnet18_20250325_114510_94.56.pth differ diff --git a/research/roi-train.py b/research/roi-train.py index 703458f..bb277ab 100755 --- a/research/roi-train.py +++ b/research/roi-train.py @@ -2,23 +2,36 @@ import torch import torch.nn as nn import torch.optim as optim -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import DataLoader, Dataset, random_split from torchvision import transforms, models from PIL import Image import os from datetime import datetime +from collections import defaultdict +import random class CustomDataset(Dataset): - def __init__(self, img_dir, transform=None): + def __init__(self, img_dir, transform=None, limit=500): self.img_dir = img_dir self.transform = transform - self.img_labels = self._load_labels() + self.img_labels = self._load_labels(limit) - def _load_labels(self): - # 假设标签存储在labels.txt文件中,每行格式为:图片名 标签 + def _load_labels(self, limit): + cats = defaultdict(list) with open(os.path.join(self.img_dir, 'labels.txt'), 'r') as f: lines = f.readlines() - return [line.strip().split() for line in lines if line.strip()] + for line in lines: + if not line.strip(): + continue + img_name, label = line.strip().split() + cats[label].append([img_name, label]) + min_samples = min(len(v) for v in cats.values()) + min_samples = min(limit, min_samples) + ret = [] + for k, v in cats.items(): + ret.extend(random.sample(v, min_samples)) + #ret.extend(v) + return ret def __len__(self): return len(self.img_labels) @@ -34,26 +47,33 @@ class CustomDataset(Dataset): # 数据预处理 transform_train = transforms.Compose([ - transforms.RandomResizedCrop((128, 64)), # 随机裁剪 - transforms.RandomHorizontalFlip(), # 随机水平翻转 + #transforms.RandomResizedCrop((128, 64)), # 随机裁剪 + #transforms.RandomHorizontalFlip(), # 随机水平翻转 + #transforms.Resize((256, 128)), # 调整大小 transforms.ToTensor(), # 转换为Tensor transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 归一化 ]) transform_val = transforms.Compose([ - transforms.Resize((128, 64)), # 调整大小 + #transforms.Resize((256, 128)), # 调整大小 transforms.ToTensor(), # 转换为Tensor transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 归一化 ]) # 加载数据集 img_dir = os.path.abspath("data/roi/train") -train_dataset = CustomDataset(img_dir=img_dir, transform=transform_train) -val_dataset = CustomDataset(img_dir=img_dir, transform=transform_val) -# 创建DataLoader +full_dataset = CustomDataset(img_dir=img_dir, transform=transform_train, limit=5000) +train_ratio = 0.5 +val_ratio = 0.5 + +train_size = int(train_ratio * len(full_dataset)) +val_size = len(full_dataset) - train_size + +train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size]) + train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) -val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False) +val_loader = DataLoader(val_dataset, batch_size=32, shuffle=True) # 加载预训练的ResNet18模型 model = models.resnet18(pretrained=True) @@ -67,12 +87,19 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) criterion = nn.CrossEntropyLoss() -optimizer = optim.Adam(model.parameters(), lr=0.001) +optimizer = optim.Adam(model.parameters(), lr=0.0001) +#optimizer = torch.optim.SGD( model.parameters(), lr=0.0001, momentum=0.09, weight_decay=1e-4) -num_epochs = 10 +scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1) + +num_epochs = 15 + +last_accu = 0 +prev_accu = [] for epoch in range(num_epochs): # 训练阶段 + print(f"Start training epoch {epoch+1}/{num_epochs}") model.train() running_loss = 0.0 for images, labels in train_loader: @@ -86,6 +113,8 @@ for epoch in range(num_epochs): print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}') + scheduler.step() + # 验证阶段 model.eval() correct = 0 @@ -98,7 +127,16 @@ for epoch in range(num_epochs): total += labels.size(0) correct += (predicted == labels).sum().item() - print(f'Validation Accuracy: {100 * correct / total:.2f}%') + last_accu = 100 * correct / total + prev_accu.append(last_accu) + if epoch > 5: + avg = sum(prev_accu[-5:]) / 5 + variance = sum((x - avg) ** 2 for x in prev_accu[-5:]) / 5 + print(f"variance={variance:.4f}") + if variance < 1 and not (prev_accu[-1] > prev_accu[-2] and prev_accu[-2] > prev_accu[-3] and prev_accu[-3] > prev_accu[-4]): + print(f"Early stopping condition met: {variance:.4f}") + break + print(f'Validation Accuracy: {last_accu:.2f}%') dt = datetime.now().strftime("%Y%m%d_%H%M%S") -torch.save(model.state_dict(), f'data/roi/resnet18_{dt}.pth') +torch.save(model.state_dict(), f'data/roi/models/resnet18_{dt}_{last_accu:.2f}.pth') diff --git a/research/roi-verify.py b/research/roi-verify.py new file mode 100755 index 0000000..19384f6 --- /dev/null +++ b/research/roi-verify.py @@ -0,0 +1,30 @@ +#!/usr/bin/env python3 +from PIL import Image +import os +import argparse +import sys +from roi_lib import * + +def parse_args(): + parser = argparse.ArgumentParser(description='ROI prediction') + parser.add_argument('--model', type=str, required=True, help='model path') + parser.add_argument('--image', type=str, required=True, help='image file') + return parser.parse_args() + +# 主函数 +def main(): + args = parse_args() + model = load_model(args.model) + image_path = args.image + image_tensor = preprocess_image(image_path) + predicted_class, probabilities = predict(model, image_tensor) + print(f'{image_path} predicted={predicted_class} prob={probabilities}') + if predicted_class == 1: + print("verify ok") + return 0 + else: + print("verify ng") + return 1 + +if __name__ == '__main__': + sys.exit(main()) \ No newline at end of file diff --git a/research/roi.py b/research/roi.py index 0f8efa4..b6596e5 100755 --- a/research/roi.py +++ b/research/roi.py @@ -9,6 +9,7 @@ import json import subprocess import cv2 from collections import defaultdict +from multiprocessing import Pool class RoiResearch(object): def __init__(self, token=None): @@ -204,22 +205,28 @@ def parse_args(): parser.add_argument("--id", "-i", type=int, action='append') return parser.parse_args() -def get_all_samples(): +def all_sample_ids(): return os.listdir("data/roi/samples") def prepare_to_train(id): rd = f"data/roi/samples/{id}" - frame_file = os.path.abspath(f"{rd}/{id}-frame.jpg") - roi_file = os.path.abspath(f"{rd}/{id}-roi.jpg") - roi_img = cv2.imread(roi_file) - frame_roi_img = frame_roi(frame_file) - frame_roi_img = cv2.resize(frame_roi_img, (128, 128)) - roi_img = cv2.resize(roi_img, (128, 128)) - side_by_side = np.concatenate((frame_roi_img, roi_img), axis=1) - side_by_side_file = os.path.abspath(f"{rd}/{id}-side-by-side.jpg") - # show_img(side_by_side, "side_by_side") - # cv2.waitKey(0) - label_file = os.path.abspath(f"{rd}/label.txt") + side_by_side_file = os.path.join(rd, f"{id}-side-by-side.jpg") + if not os.path.exists(side_by_side_file): + frame_file = os.path.abspath(f"{rd}/{id}-frame.jpg") + roi_file = os.path.abspath(f"{rd}/{id}-roi.jpg") + if not os.path.exists(frame_file) or not os.path.exists(roi_file): + print(f"skipping {id}, no frame or roi") + return + roi_img = cv2.imread(roi_file) + try: + frame_roi_img = frame_roi(frame_file) + except Exception as e: + print(f"failed to get frame_roi for {id}: {e}") + return + frame_roi_img = cv2.resize(frame_roi_img, (128, 128)) + roi_img = cv2.resize(roi_img, (128, 128)) + side_by_side = np.concatenate((frame_roi_img, roi_img), axis=1) + cv2.imwrite(side_by_side_file, side_by_side) json_file = os.path.abspath(f"{rd}/{id}.json") with open(json_file, "r") as f: data = json.load(f) @@ -230,9 +237,9 @@ def prepare_to_train(id): elif 'neg' in labels: label = 0 else: - raise Exception("no label found") - side_by_side_file = os.path.abspath(f"data/roi/train/{id}.jpg") - cv2.imwrite(side_by_side_file, side_by_side) + print(f"no label found for {id}") + return + shutil.copy(side_by_side_file, os.path.abspath(f"data/roi/train/{id}.jpg")) with open(os.path.abspath(f"data/roi/train/labels.txt"), "a") as f: f.write(f"{id}.jpg {label}\n") return label @@ -242,19 +249,10 @@ def main(): if args.download: get_roi_data(args) if args.preprocess: - all_samples = get_all_samples() - total = len(all_samples) shutil.rmtree(os.path.abspath("data/roi/train")) os.makedirs(os.path.abspath("data/roi/train")) - label_count = defaultdict(int) - for i, id in enumerate(all_samples): - print(f"preprocessing {id} ({i + 1}/{total})") - try: - label = prepare_to_train(id) - label_count[label] += 1 - except Exception as e: - print(f"error: {e}") - print(f"count by label: {label_count}") + with Pool(processes=10) as pool: + pool.map(prepare_to_train, all_sample_ids()) return if args.id: for id in args.id: diff --git a/research/roi_bench.py b/research/roi_bench.py new file mode 100755 index 0000000..7ec93df --- /dev/null +++ b/research/roi_bench.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +from PIL import Image +import os +import argparse +import random +from roi_lib import * + + +def parse_args(): + parser = argparse.ArgumentParser(description='ROI prediction') + parser.add_argument('--model', type=str, required=True, help='model path') + parser.add_argument('--label', type=str, required=True, help='label file') + return parser.parse_args() + +# 主函数 +def main(): + args = parse_args() + # 加载模型 + model = load_model(args.model) + with open(args.label, 'r') as f: + lines = f.readlines() + correct = 0 + zero_to_one = 0 + one_to_zero = 0 + testset = [] + for line in lines: + if not line.strip(): + continue + fs = line.strip().split() + if len(fs) != 2: + continue + image_name = fs[0] + label = fs[1] + testset.append((image_name, label)) + random.shuffle(testset) + done = 0 + for image_name, label in testset: + image_path = os.path.abspath(os.path.join(os.path.dirname(args.label), image_name)) + image_tensor = preprocess_image(image_path) + predicted_class, probabilities = predict(model, image_tensor) + done += 1 + if str(predicted_class) == label: + correct += 1 + else: + if label == '0': + zero_to_one += 1 + else: + one_to_zero += 1 + accu_so_far = correct / done + print(f'{image_path} label={label} predicted={predicted_class} prob={probabilities} accuracy so far: {accu_so_far:.3f}') + print(f'total={done} correct={correct} accuracy={accu_so_far:.3f}') + print(f'zero_to_one={zero_to_one} one_to_zero={one_to_zero}') + return 0 + +if __name__ == '__main__': + sys.exit(main()) diff --git a/research/roi_lib.py b/research/roi_lib.py new file mode 100755 index 0000000..d360d8a --- /dev/null +++ b/research/roi_lib.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +import torch +import torch.nn as nn +import torchvision.models as models +import torchvision.transforms as transforms +from PIL import Image +import os +import argparse +import random + +def load_model(model_path): + model = models.resnet18(pretrained=True) + num_ftrs = model.fc.in_features + model.fc = nn.Linear(num_ftrs, 2) + model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) # 加载模型权重 + model.eval() + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + model = model.to(device) + return model + +def preprocess_image(image_path): + transform = transforms.Compose([ + transforms.ToTensor(), # 转换为Tensor + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 归一化 + ]) + image = Image.open(image_path).convert('RGB') # 打开图像并转换为 RGB + image = transform(image).unsqueeze(0) # 增加 batch 维度 + return image + +def predict(model, image_tensor): + with torch.no_grad(): # 禁用梯度计算 + output = model(image_tensor) + _, predicted = torch.max(output, 1) # 获取预测类别 + probabilities = torch.nn.functional.softmax(output, dim=1) # 计算概率 + return predicted.item(), probabilities.squeeze().tolist() \ No newline at end of file diff --git a/research/server.py b/research/server.py index 29e5d16..d24dcd0 100755 --- a/research/server.py +++ b/research/server.py @@ -11,6 +11,7 @@ from io import BytesIO import tempfile import cv2 import numpy as np +from roi_lib import * app = Flask(__name__) @@ -236,5 +237,20 @@ def analyze(): analyze_datapoint(data) return data +@app.route('/api/roi-verify', methods=['POST']) +def roi_verify(): + files = request.files + side_by_side_fn = files['side_by_side'] + with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as f: + side_by_side_fn.save(f.name) + img = Image.open(f.name) + img = preprocess_image(img) + predicted_class, probabilities = predict(model, image_tensor) + return { + "ok": True, + "predicted_class": predicted_class, + "probabilities": probabilities, + } + if __name__ == '__main__': app.run(host='0.0.0.0', port=26966, debug=True)