简介：这份数据集面向计算机视觉入门者、深度学习教学者以及需要轻量级分类实验的开发者，提供星形、圆形、正方形、三角形四类基本几何形状的标注图像，可用于图像分类模型的训练、验证与教学演示，帮助快速搭建形状识别实验环境。资源包共收录14973个文件，其中14970个为png格式图片，另有3个py脚本，压缩包整体约21.57MB，图片统一为200×200像素，尺寸规整、类别清晰，便于直接读取与批量预处理。目前已有293人学习下载，适合作为课程作业、模型对比或数据增强练习的素材。读者可获得一套开箱即用的四分类图像数据，配合附带的Python脚本可快速完成数据加载与目录遍历，省去自行采集与清洗的成本，也能在此基础上尝试迁移学习、卷积网络调参等进阶任务，是入门图像分类的实用起点。
1. 四种基本形状图片数据集：从零训练一个形状分类器到底值不值
手头有个小需求：给流水线上的零件做初筛，只判断它是星形、圆形、正方形还是三角形。你不需要 ImageNet 级别的模型，也不需要 GPU 集群，一张 224×224 的小图、一个四分类头，几秒钟就能出结果。问题在于——数据从哪来？网上现成的形状数据集要么是手写数字那种灰度图，要么是几何图形合成图，类别不齐、分辨率参差，拿来做真实场景的迁移几乎必翻车。四种基本形状的图片数据集（星形、圆形、正方形、三角形）就是为这类场景准备的：类别少、语义清晰、可控性强，适合做教学演示、原型验证、边缘设备上的轻量分类。这篇文章不讲空泛的“数据集很重要”，而是把从拿到数据到跑通推理的整条链路拆开，告诉你每一步的参数怎么设、坑在哪、什么情况下这个方向不值得投入。如果你正在做形状识别、几何特征分类、或者想找一个干净的小数据集练手，下面的内容可以直接抄作业。
2. 形状数据集怎么选、怎么造：合成与实拍的取舍
2.1 四种形状的数据集长什么样才算“能用”
一个能直接训练的形状数据集，核心不是图片数量，而是类间差异是否被噪声淹没。星形、圆形、正方形、三角形这四类，在理想几何条件下是线性可分的，但一旦加入旋转、缩放、光照、背景纹理，边界就模糊了。我一般会先检查三个指标：每类样本数是否均衡（偏差不超过 1.2:1）、图像分辨率是否统一（至少短边 128px）、背景是否单一。如果背景里混入了文字、线条、其他形状的残影，模型会学到“背景纹理”而不是“形状轮廓”，这是最常见的翻车点。
合成数据集的优势在于标注零成本、类别绝对均衡。用 OpenCV 或 PIL 画一个星形，顶点坐标固定，旋转角度可控，背景纯白或纯黑，生成一万张也就几十秒。但合成数据的致命伤是域偏移：真实摄像头拍出来的形状有透视变形、边缘模糊、反光，合成图里没有。我的做法是合成打底、实拍微调——先用 8000 张合成图训练一个基础模型，再用 200 张实拍图做最后一层微调，准确率能从 0.82 拉到 0.94 左右。
实拍数据集则相反，标注成本高，但泛化好。如果你只有手机能拍，建议在纯色卡纸上剪出四种形状，固定光源、固定距离，每类拍 300 张，旋转角度覆盖 0° 到 360°。注意：不要只拍正面，倾斜 15° 到 30° 的样本必须占至少 20%，否则模型对透视变换毫无抵抗力。
2.2 用 Python 合成四种形状图片的最小脚本
下面这段代码用 PIL 生成星形、圆形、正方形、三角形各 500 张，随机旋转、随机缩放、随机背景色，输出到四个子目录。依赖只有 Pillow 和 numpy。
import os import math import random from PIL import Image, ImageDraw # 输出根目录，每个类别一个子文件夹 ROOT = "shapes_dataset" CLASSES = ["star", "circle", "square", "triangle"] IMG_SIZE = 224 SAMPLES_PER_CLASS = 500 # 随机背景色范围：浅灰到白，避免纯黑导致边缘对比过强 BG_RANGE = (200, 255) # 形状填充色范围：深色，与背景形成对比 FG_RANGE = (0, 80) def random_color(low, high): return tuple(random.randint(low, high) for _ in range(3)) def draw_star(draw, cx, cy, r, fill): # 五角星：外顶点半径 r，内顶点半径 r*0.4 points = [] for i in range(10): angle = math.pi / 2 + i * math.pi / 5 radius = r if i % 2 == 0 else r * 0.4 x = cx + radius * math.cos(angle) y = cy - radius * math.sin(angle) points.append((x, y)) draw.polygon(points, fill=fill) def draw_triangle(draw, cx, cy, r, fill): # 等边三角形，顶点朝上 points = [] for i in range(3): angle = math.pi / 2 + i * 2 * math.pi / 3 x = cx + r * math.cos(angle) y = cy - r * math.sin(angle) points.append((x, y)) draw.polygon(points, fill=fill) def generate_one(cls, idx): img = Image.new("RGB", (IMG_SIZE, IMG_SIZE), random_color(*BG_RANGE)) draw = ImageDraw.Draw(img) # 中心点随机偏移 ±20px，防止模型只学“居中”特征 cx = IMG_SIZE // 2 + random.randint(-20, 20) cy = IMG_SIZE // 2 + random.randint(-20, 20) # 半径随机 40 到 80，控制形状大小 r = random.randint(40, 80) fill = random_color(*FG_RANGE) if cls == "star": draw_star(draw, cx, cy, r, fill) elif cls == "circle": draw.ellipse([cx - r, cy - r, cx + r, cy + r], fill=fill) elif cls == "square": draw.rectangle([cx - r, cy - r, cx + r, cy + r], fill=fill) elif cls == "triangle": draw_triangle(draw, cx, cy, r, fill) # 随机旋转 -30 到 30 度，模拟拍摄角度偏差 angle = random.uniform(-30, 30) img = img.rotate(angle, resample=Image.BILINEAR, expand=False) return img for cls in CLASSES: out_dir = os.path.join(ROOT, cls) os.makedirs(out_dir, exist_ok=True) for i in range(SAMPLES_PER_CLASS): img = generate_one(cls, i) img.save(os.path.join(out_dir, f"{cls}_{i:04d}.png")) print("done")逻辑说明：draw_star用极坐标生成十个顶点，外顶点半径r、内顶点r*0.4，这是五角星的标准画法。draw_triangle同理，三个顶点均匀分布在圆周上。中心点偏移和随机旋转是为了让模型不依赖绝对位置和角度。BG_RANGE和FG_RANGE的取值决定了对比度，如果背景太暗、形状太亮，边缘检测会失效，所以我把背景限制在浅色区间。
参数说明：IMG_SIZE改成 128 可以加快训练但会损失细节；SAMPLES_PER_CLASS建议不低于 300，否则四分类的验证集波动会很大；r的范围 40 到 80 对应形状占图比例约 18% 到 36%，太小则形状像素不足，太大则旋转后容易出界。生成完检查一下每个文件夹的图片数量是否一致，不一致说明某类生成时抛了异常被吞掉。
2.3 实拍补充：手机拍 200 张的采集规范
合成数据跑通后，如果目标场景是真实摄像头，必须补实拍。我一般用手机固定在三脚架上，俯拍一张 A4 纸，纸上贴四种形状的卡纸。每换一个形状，旋转 8 个角度（0°、45°、90°……315°），每个角度拍 6 张，加上不同光照（顺光、侧光、背光）各 2 张，总共 4×8×6=192 张，接近 200 张。拍完用 LabelImg 或 Label Studio 标注，但四分类任务其实不需要画框，直接按文件夹分好即可。注意：实拍图要统一裁剪到 224×224，裁剪时保证形状居中，边缘留 20px 余量。如果实拍图里出现了合成数据没有的阴影，不要试图用后期去掉，保留阴影反而能提升模型鲁棒性。
3. 用 PyTorch 训练四分类模型：从 DataLoader 到混淆矩阵
3.1 数据划分与增强策略的参数怎么定
拿到shapes_dataset后，按 7:2:1 划分训练集、验证集、测试集。注意：如果合成图和实拍图混在一起，必须保证同一张原图的不同增强版本不会跨集出现，否则验证集准确率会虚高。我一般先用torchvision.datasets.ImageFolder加载，再用random_split按类别分层划分。增强策略只保留三种：随机水平翻转、随机旋转 ±15°、颜色抖动（亮度 0.2、对比度 0.2、饱和度 0.2）。不要加随机裁剪，因为形状本身占图比例不大，裁多了会把形状裁掉。
import torch from torch.utils.data import DataLoader, random_split from torchvision import datasets, transforms # 训练增强：翻转、旋转、颜色抖动 train_tf = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(degrees=15), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # 验证/测试：只做 Resize 和归一化 val_tf = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) full_dataset = datasets.ImageFolder("shapes_dataset", transform=train_tf) # 按 7:2:1 划分，固定随机种子保证可复现 torch.manual_seed(42) n_total = len(full_dataset) n_train = int(0.7 * n_total) n_val = int(0.2 * n_total) n_test = n_total - n_train - n_val train_set, val_set, test_set = random_split(full_dataset, [n_train, n_val, n_test]) # 验证集要换成 val_tf，这里用 Subset 包装时需重新设置 transform # 简单做法：分别创建 dataset 再取索引，或者用 DataLoader 的 collate 里不做增强 train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=2) val_loader = DataLoader(val_set, batch_size=32, shuffle=False, num_workers=2) test_loader = DataLoader(test_set, batch_size=32, shuffle=False, num_workers=2)逻辑说明：random_split直接按长度切分，不保证类别均衡，所以最好先统计每类数量再手动分层。Normalize的 mean 和 std 都设 0.5 是因为合成图背景接近白色，实拍图也以浅色为主，用 ImageNet 的均值反而会引入偏差。batch_size=32在 224×224 输入下显存占用约 1.2GB，6GB 显存的卡就能跑。
参数说明：RandomRotation(degrees=15)比 30 更安全，因为训练集里已经有合成时的 ±30° 旋转，再叠加容易导致形状出界。ColorJitter的三个系数不要超过 0.3，否则颜色失真会让形状边缘和背景混在一起。num_workers在 Windows 上设 0，Linux 上设 2 到 4，设太大反而拖慢小数据集的加载。
3.2 一个四分类 CNN 的最小训练循环
模型不用 ResNet，自己搭一个 4 层卷积加全局平均池化就够了。参数量控制在 50 万以内，推理时 CPU 单张图 8ms 左右。
import torch.nn as nn import torch.optim as optim class ShapeNet(nn.Module): def __init__(self, num_classes=4): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1) ) self.classifier = nn.Linear(64, num_classes) def forward(self, x): x = self.features(x) x = x.view(x.size(0), -1) return self.classifier(x) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = ShapeNet(num_classes=4).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=1e-3) for epoch in range(30): model.train() running_loss = 0.0 for imgs, labels in train_loader: imgs, labels = imgs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(imgs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() # 验证阶段 model.eval() correct, total = 0, 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels = imgs.to(device), labels.to(device) preds = model(imgs).argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) print(f"epoch {epoch+1}, loss {running_loss/len(train_loader):.4f}, val_acc {correct/total:.4f}")逻辑说明：AdaptiveAvgPool2d(1)把 28×28 的特征图压成 1×1，替代全连接层，减少参数量。CrossEntropyLoss内部带 softmax，不需要手动加。学习率 1e-3 配合 Adam 在 30 个 epoch 内基本收敛，如果验证集准确率在第 15 个 epoch 后不再上升，可以降到 1e-4 再跑 10 个 epoch。
参数说明：lr=1e-3是 Adam 的常用起点，如果 loss 震荡剧烈，改成 5e-4。epoch=30是针对 2000 张训练图的经验值，图多了要加。验证时记得model.eval()和torch.no_grad()，否则 BatchNorm 和 Dropout 会干扰结果。如果显存不够，把batch_size降到 16，同时学习率乘以 0.5。
3.3 混淆矩阵告诉你哪两类在互相误判
训练完在测试集上跑一遍，输出混淆矩阵。四分类的混淆矩阵是 4×4，行是真实标签，列是预测标签。我见过最多的误判是星形和三角形——五角星的尖角在低分辨率下和三角形的顶点很像。如果星形被误判为三角形的比例超过 10%，说明模型没有学到“内凹”特征。解决办法：在训练集里增加星形的内凹角度变化，或者把输入分辨率从 224 提到 256。
from sklearn.metrics import confusion_matrix import numpy as np model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for imgs, labels in test_loader: imgs = imgs.to(device) preds = model(imgs).argmax(dim=1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) cm = confusion_matrix(all_labels, all_preds) print(cm) # 输出示例： # [[48 1 0 1] # [ 0 50 0 0] # [ 0 0 49 1] # [ 2 0 0 48]]逻辑说明：confusion_matrix的第一个参数是真实标签，第二个是预测标签。对角线是正确分类数，非对角线是误判。如果某一列的和远大于对角线，说明模型倾向于把其他类预测成这一类。参数说明：test_loader的shuffle必须为 False，否则标签和预测对不上。如果测试集太小（比如每类不到 30 张），混淆矩阵的波动会很大，建议至少每类 50 张。
4. 形状分类的避坑与排查：那些让准确率卡在 0.7 的细节
4.1 现象：验证集准确率 0.99，测试集掉到 0.6
原因：合成图和实拍图混在一起划分时，同一张原图的增强版本跨了集。比如原图 A 的旋转版在训练集，原图 A 的翻转版在测试集，模型见过“类似”的图，测试集自然虚高。解决：按原图 ID 划分，同一 ID 的所有增强版本只进一个集。如果合成图和实拍图都有，先按来源分层，再按 ID 划分。
4.2 现象：模型把圆形和正方形都预测成三角形
原因：背景色和形状填充色对比度太低，边缘模糊，卷积核学不到有效梯度。检查一下BG_RANGE和FG_RANGE是否有重叠。解决：把背景限制在 220 到 255，填充限制在 0 到 60，确保灰度差至少 160。实拍图如果背景是灰色桌面，换一张白纸垫底。
4.3 现象：训练 loss 降到 0.01 后不再下降，但验证 loss 开始上升
原因：过拟合。四分类任务参数量超过 100 万就容易过拟合，尤其是合成图纹理单一。解决：把ShapeNet的最后一层卷积通道从 64 降到 32，或者加nn.Dropout(0.3)在分类头之前。另外，早停策略设patience=5，验证 loss 连续 5 个 epoch 不降就停。
4.4 现象：推理时单张图预测结果和批量预测不一致
原因：model.eval()没加，或者输入归一化的 mean/std 和训练时不一致。检查推理代码里是否用了val_tf而不是train_tf。解决：把预处理封装成一个函数，训练和推理共用，避免手写两套。另外，ToTensor()会把像素值从 0-255 缩到 0-1，如果推理时忘了这一步，输入范围变成 0-255，模型输出会完全乱掉。
4.5 现象：星形和三角形混淆严重，调参无效
原因：星形的内凹特征在 224×224 下只有几个像素宽，卷积核感受野不够。解决：把输入分辨率提到 320×320，或者在第一层卷积后用nn.MaxPool2d(2)之前加一个nn.Conv2d(3, 16, 5, padding=2)增大感受野。另一个办法是数据层面：生成星形时把内凹半径从r*0.4改成r*0.3，让内凹更明显。
5. 把模型塞进边缘设备：ONNX 导出与量化后的精度验证
训练完的 PyTorch 模型如果要部署到树莓派或 Jetson Nano，直接装 PyTorch 太重。我一般导出成 ONNX，再用 ONNX Runtime 做 INT8 量化。导出时注意opset_version选 11 或 13，输入动态轴只保留 batch 维度。
import torch.onnx model.eval() dummy_input = torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, "shape_net.onnx", input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}, opset_version=11 )逻辑说明：dynamic_axes让 batch 维度可变，推理时可以一次送多张图。opset_version=11兼容大多数 ONNX Runtime 版本。导出后用onnxruntime加载，对比 PyTorch 和 ONNX 的输出差异，如果最大绝对误差超过 1e-4，检查是否有不支持的自定义层。
量化用 ONNX Runtime 的quantize_static，需要准备 100 张校准图。校准集从训练集里随机抽，覆盖四类。量化后模型大小从 2MB 降到 500KB，推理速度提升约 2 倍，但精度可能掉 1 到 3 个百分点。如果掉超过 5 个点，说明校准集分布和测试集差异太大，换一批校准图。
验证量化模型时，不要只看整体准确率，要重新算一遍混淆矩阵。我遇到过量化后圆形召回率从 0.98 掉到 0.85 的情况，原因是圆形的边缘像素在 INT8 下被截断。解决办法：对圆形类做少量微调，或者把量化改成per_channel模式。
最后说个习惯：每次导出 ONNX 后，我都会用同一张测试图分别跑 PyTorch 和 ONNX，打印 top-1 类别和置信度，两者一致才继续往下走。这个检查花不了两分钟，但能省掉后面几个小时的排查。希望帮到你。
本文还有配套的精品资源，点击获取
