research: Add roi-train.sh

research/roi-train.py

@@ -9,28 +9,23 @@ import os
 from datetime import datetime
 from collections import defaultdict
 import random
+import argparse
 
-class CustomDataset(Dataset):
+class SideBySideDataset(Dataset):
-    def __init__(self, img_dir, transform=None, limit=500):
+    def __init__(self, img_dir, labels_file, transform=None):
         self.img_dir = img_dir
         self.transform = transform
-        self.img_labels = self._load_labels(limit)
+        self.img_labels = self._load_labels(labels_file)
 
-    def _load_labels(self, limit):
+    def _load_labels(self, labels_file):
-        cats = defaultdict(list)
+        ret = []
-        with open(os.path.join(self.img_dir, 'labels.txt'), 'r') as f:
+        with open(labels_file, 'r') as f:
             lines = f.readlines()
         for line in lines:
             if not line.strip():
                 continue
             img_name, label = line.strip().split()
-            cats[label].append([img_name, label])
+            ret.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):
@@ -45,98 +40,112 @@ class CustomDataset(Dataset):
         label = int(label)
         return image, label
 
-# 数据预处理
+def parse_args():
-transform_train = transforms.Compose([
+    parser = argparse.ArgumentParser(description='Train a model')
-    #transforms.RandomResizedCrop((128, 64)),  # 随机裁剪
+    parser.add_argument('--labels-file', required=True, type=str, help='Path to the labels file')
-    #transforms.RandomHorizontalFlip(),       # 随机水平翻转
+    return parser.parse_args()
-    #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([
+def main():
-    #transforms.Resize((256, 128)),            # 调整大小
+    args = parse_args()
-    transforms.ToTensor(),                   # 转换为Tensor
+    # 数据预处理
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
+    transform_train = transforms.Compose([
-])
+        #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([
-img_dir = os.path.abspath("data/roi/train")
+        #transforms.Resize((256, 128)),            # 调整大小
+        transforms.ToTensor(),                   # 转换为Tensor
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
+    ])
 
-full_dataset = CustomDataset(img_dir=img_dir, transform=transform_train, limit=5000)
+    # 加载数据集
-train_ratio = 0.5
+    img_dir = os.path.abspath("data/roi/train")
-val_ratio = 0.5
 
-train_size = int(train_ratio * len(full_dataset))
+    full_dataset = SideBySideDataset(
-val_size = len(full_dataset) - train_size
+        img_dir=img_dir,
+        labels_file=args.labels_file,
+        transform=transform_train,
+    )
+    train_ratio = 0.5
+    val_ratio = 0.5
 
-train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])
+    train_size = int(train_ratio * len(full_dataset))
+    val_size = len(full_dataset) - train_size
 
-train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+    train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size])
-val_loader = DataLoader(val_dataset, batch_size=32, shuffle=True)
 
-# 加载预训练的ResNet18模型
+    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
-model = models.resnet18(pretrained=True)
+    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=True)
 
-# 修改最后一层全连接层，使其输出为2（二分类）
+    # 加载预训练的ResNet18模型
-num_ftrs = model.fc.in_features
+    model = models.resnet18(pretrained=True)
-model.fc = nn.Linear(num_ftrs, 2)
 
-# 将模型移动到GPU（如果可用）
+    # 修改最后一层全连接层，使其输出为2（二分类）
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    num_ftrs = model.fc.in_features
-model = model.to(device)
+    model.fc = nn.Linear(num_ftrs, 2)
 
-criterion = nn.CrossEntropyLoss()
+    # 将模型移动到GPU（如果可用）
-optimizer = optim.Adam(model.parameters(), lr=0.0001)
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-#optimizer = torch.optim.SGD( model.parameters(), lr=0.0001, momentum=0.09, weight_decay=1e-4)
+    model = model.to(device)
 
-scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
+    criterion = nn.CrossEntropyLoss()
+    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 = 15
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
 
-last_accu = 0
+    num_epochs = 15
-prev_accu = []
 
-for epoch in range(num_epochs):
+    last_accu = 0
-    # 训练阶段
+    prev_accu = []
-    print(f"Start training epoch {epoch+1}/{num_epochs}")
-    model.train()
-    running_loss = 0.0
-    for images, labels in train_loader:
-        images, labels = images.to(device), labels.to(device)
-        optimizer.zero_grad()
-        outputs = model(images)
-        loss = criterion(outputs, labels)
-        loss.backward()
-        optimizer.step()
-        running_loss += loss.item()
 
-    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')
+    for epoch in range(num_epochs):
-
+        # 训练阶段
-    scheduler.step()
+        print(f"Start training epoch {epoch+1}/{num_epochs}")
-
+        model.train()
-    # 验证阶段
+        running_loss = 0.0
-    model.eval()
+        for images, labels in train_loader:
-    correct = 0
-    total = 0
-    with torch.no_grad():
-        for images, labels in val_loader:
             images, labels = images.to(device), labels.to(device)
+            optimizer.zero_grad()
             outputs = model(images)
-            _, predicted = torch.max(outputs.data, 1)
+            loss = criterion(outputs, labels)
-            total += labels.size(0)
+            loss.backward()
-            correct += (predicted == labels).sum().item()
+            optimizer.step()
+            running_loss += loss.item()
 
-    last_accu = 100 * correct / total
+        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')
-    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")
+        scheduler.step()
-torch.save(model.state_dict(), f'data/roi/models/resnet18_{dt}_{last_accu:.2f}.pth')
+
+        # 验证阶段
+        model.eval()
+        correct = 0
+        total = 0
+        with torch.no_grad():
+            for images, labels in val_loader:
+                images, labels = images.to(device), labels.to(device)
+                outputs = model(images)
+                _, predicted = torch.max(outputs.data, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+
+        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/models/resnet18_{dt}_{last_accu:.2f}.pth')
+
+if __name__ == "__main__":
+    main()

research/roi-train.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+all_label_file="data/roi/train/labels.txt"
+npos=$(grep '1$' $all_label_file | wc -l)
+nneg=$(grep '0$' $all_label_file | wc -l)
+
+echo "npos: $npos"
+echo "nneg: $nneg"
+
+min=$npos
+if [ $nneg -lt $min ]; then
+    min=$nneg
+fi
+
+echo "min: $min"
+cat $all_label_file | grep '1$' | shuf | head -n $min > data/roi/train/pos.txt
+cat $all_label_file | grep '0$' | shuf | head -n $min > data/roi/train/neg.txt
+
+cat data/roi/train/pos.txt data/roi/train/neg.txt > data/roi/train/mixed.txt
+
+# these are the 'original' prints with different bg pattern offset
+grep -e "7792[1-7].jpg" $all_label_file >> data/roi/train/mixed.txt
+
+./roi-train.py --labels-file data/roi/train/mixed.txt

research/roi-verify.py

@@ -8,23 +8,24 @@ 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')
+    parser.add_argument('image', nargs='+', type=str, help='image file')
     return parser.parse_args()
 
 # 主函数
 def main():
     args = parse_args()
     model = load_model(args.model)
-    image_path = args.image
+    ret = 0
-    image_tensor = preprocess_image(image_path)
+    for image_path in args.image:
-    predicted_class, probabilities = predict(model, image_tensor)
+        image_tensor = preprocess_image(image_path)
-    print(f'{image_path} predicted={predicted_class} prob={probabilities}')
+        predicted_class, probabilities = predict(model, image_tensor)
-    if predicted_class == 1:
+        print(f'{image_path} predicted={predicted_class} prob={probabilities}')
-        print("verify ok")
+        if predicted_class == 1:
-        return 0
+            print("verify ok")
-    else:
+        else:
-        print("verify ng")
+            print("verify ng")
-        return 1
+            ret = 1
+    return ret
 
 if __name__ == '__main__':
     sys.exit(main())