MiniGPT-4/data/create_mvtec_dataset.ipynb

385 lines
124 KiB
Plaintext
Raw Normal View History

2025-01-12 00:43:42 +00:00
{
"cells": [
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 6,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from pathlib import Path\n",
"import glob\n",
"import cv2\n",
"import json"
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 7,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [],
"source": [
"BASE_DIR = Path(\"./MVTEC\")\n",
"NEW_DIR = Path(\"./MVTEC_det\")\n",
"NEW_DIR.mkdir(exist_ok=True)\n",
"\n",
"class_names = [\"bottle\", \"cable\", \"capsule\", \"carpet\", \"grid\", \"hazelnut\", \"leather\", \"metal_nut\", \"pill\", \"screw\", \"tile\", \"toothbrush\", \"transistor\", \"wood\", \"zipper\"]\n",
"base_split = [\"train\", \"test\"]\n",
"base_cate = [\"good\", \"defect\"]"
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"def normalize_bbox(bbox, width, height):\n",
" \"\"\"\n",
" Normalize bounding box to the range [0, 100].\n",
" \"\"\"\n",
" return [\n",
" int((bbox[0] / width) * 100), # x_min\n",
" int((bbox[1] / height) * 100), # y_min\n",
" int((bbox[2] / width) * 100), # x_max\n",
" int((bbox[3] / height) * 100), # y_max\n",
" ]"
]
},
{
"cell_type": "code",
"execution_count": 9,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Processing bottle\n",
"Processing cable\n",
"Processing capsule\n",
"Processing carpet\n",
"Processing grid\n",
"Processing hazelnut\n",
"Processing leather\n",
"Processing metal_nut\n",
"Processing pill\n",
"Processing screw\n",
"Processing tile\n",
"Processing toothbrush\n",
"Processing transistor\n",
"Processing wood\n",
"Processing zipper\n"
]
}
],
"source": [
"mvtech_ad_data_for_regression = []\n",
"\n",
"for class_name in class_names:\n",
" print(f\"Processing {class_name}\")\n",
" class_dir = Path(BASE_DIR) / class_name\n",
" new_class_dir = NEW_DIR / class_name\n",
" new_class_dir.mkdir(exist_ok=True)\n",
" \n",
" for split in base_split:\n",
" if split == \"train\":\n",
" cmd = f\"cp -r {class_dir / split}/good {new_class_dir}/\"\n",
" os.system(cmd)\n",
"\n",
" image_paths = glob.glob(f\"{new_class_dir}/good/*.png\")\n",
" for image_path in image_paths:\n",
" img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)\n",
" # bbox is the image size\n",
" bbox = [0, 0, img.shape[1], img.shape[0]]\n",
" mvtech_ad_data_for_regression.append({\n",
" \"image_path\": image_path,\n",
" \"bbox\": bbox,\n",
" \"class\": class_name,\n",
" \"is_broken\": False,\n",
" \"height\": img.shape[0],\n",
" \"width\": img.shape[1]\n",
" })\n",
"\n",
" else:\n",
" new_broken_dir = new_class_dir / \"broken\"\n",
" new_broken_dir.mkdir(exist_ok=True)\n",
"\n",
" last_good_idx = int(sorted(glob.glob(f'{class_dir}/train/good/*.png'))[-1].split('/')[-1].split('.')[0])\n",
" good_images = sorted(glob.glob(f\"{class_dir / split}/good/*.png\"))\n",
" for i, good_image in enumerate(good_images):\n",
" new_name = good_image.split(\"/\")[:-1] + [f\"{last_good_idx+i+1:03d}.png\"]\n",
" new_name.remove(\"test\")\n",
" new_name = \"/\".join(new_name)\n",
" new_name = new_name.replace(\"MVTEC\", \"MVTEC_det\")\n",
" cmd = f\"cp {good_image} {new_name}\"\n",
" os.system(cmd)\n",
"\n",
" img = cv2.imread(new_name, cv2.IMREAD_GRAYSCALE)\n",
" bbox = [0, 0, img.shape[1], img.shape[0]]\n",
" mvtech_ad_data_for_regression.append({\n",
" \"image_path\": new_name,\n",
" \"bbox\": bbox,\n",
" \"class\": class_name,\n",
" \"is_broken\": False,\n",
" \"height\": 256,\n",
" \"width\": 256\n",
" })\n",
"\n",
" broken_idx = 0\n",
" broken_cates = [x.split(\"/\")[-1] for x in glob.glob(f\"{class_dir / split}/*\")]\n",
" broken_cates.remove(\"good\")\n",
" for broken_cate in broken_cates:\n",
" image_paths = sorted(glob.glob(f\"{class_dir / split}/{broken_cate}/*.png\"))\n",
" for image_path in image_paths:\n",
" new_name = image_path.split(\"/\")[:-1] + [f\"{broken_idx:03d}.png\"]\n",
" new_name.remove(\"test\")\n",
" new_name = \"/\".join(new_name)\n",
" new_name = new_name.replace(broken_cate, \"broken\")\n",
" new_name = new_name.replace(\"MVTEC\", \"MVTEC_det\")\n",
" cmd = f\"cp {image_path} {new_name}\"\n",
" os.system(cmd)\n",
"\n",
" mask_path = image_path.replace(\"test\", \"ground_truth\").replace(\".png\", \"_mask.png\")\n",
" mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)\n",
" contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
" if not contours:\n",
" raise ValueError(\"No contours found\")\n",
" x, y, w, h = cv2.boundingRect(contours[0])\n",
" bbox = [x, y, x + w, y + h]\n",
" mvtech_ad_data_for_regression.append({\n",
" \"image_path\": new_name,\n",
" \"bbox\": bbox,\n",
" \"class\": class_name,\n",
" \"is_broken\": True,\n",
" \"height\": mask.shape[0],\n",
" \"width\": mask.shape[1]\n",
" })\n",
"\n",
" broken_idx += 1\n",
" "
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 10,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"5354"
]
},
2025-01-13 04:07:02 +00:00
"execution_count": 10,
2025-01-12 00:43:42 +00:00
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"len(mvtech_ad_data_for_regression)"
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 11,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [],
"source": [
"# shuffle the data\n",
"\n",
"import random\n",
"random.seed(0)\n",
"\n",
"random.shuffle(mvtech_ad_data_for_regression)"
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 12,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [],
"source": [
"for data in mvtech_ad_data_for_regression:\n",
2025-01-13 04:07:02 +00:00
" data[\"image_path\"] = data[\"image_path\"].replace(\"MVTEC_det\", \".\")\n",
" data[\"bbox\"] = normalize_bbox(\n",
" bbox=data[\"bbox\"],\n",
" width=data[\"width\"],\n",
" height=data[\"height\"]\n",
" )"
2025-01-12 00:43:42 +00:00
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 13,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [],
"source": [
"with open(\"./MVTEC_det/mvtech_ad_data_for_regression.json\", \"w\") as f:\n",
" json.dump(mvtech_ad_data_for_regression, f, indent=4)"
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 14,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [],
"source": [
"# Move all the images to the images folder\n",
"\n",
"os.makedirs(\"./MVTEC_det/images\", exist_ok=True)\n",
"\n",
"for class_name in class_names:\n",
" cmd = f\"mv MVTEC_det/{class_name} MVTEC_det/images/\"\n",
" os.system(cmd)"
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 15,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [],
"source": [
"# split data into train test val\n",
"train_ratio = 0.7\n",
"val_ratio = 0.1\n",
"test_ratio = 0.2\n",
"\n",
"train_data = mvtech_ad_data_for_regression[:int(len(mvtech_ad_data_for_regression) * train_ratio)]\n",
"val_data = mvtech_ad_data_for_regression[int(len(mvtech_ad_data_for_regression) * train_ratio):int(len(mvtech_ad_data_for_regression) * (train_ratio + val_ratio))]\n",
"test_data = mvtech_ad_data_for_regression[int(len(mvtech_ad_data_for_regression) * (train_ratio + val_ratio)):]\n",
"\n",
"\n",
"\n",
"with open(\"./MVTEC_det/train_data.json\", \"w\") as f:\n",
" json.dump(train_data, f, indent=4)\n",
"\n",
"with open(\"./MVTEC_det/val_data.json\", \"w\") as f:\n",
" json.dump(val_data, f, indent=4)\n",
"\n",
"with open(\"./MVTEC_det/test_data.json\", \"w\") as f:\n",
" json.dump(test_data, f, indent=4)"
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 16,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [],
"source": [
"# Statistic the data balance of the train, val, test set between classes and between good and defect\n",
"\n",
"train_class_count = {class_name: 0 for class_name in class_names}\n",
"train_good_defect_count = {\"good\": 0, \"defect\": 0}\n",
"\n",
"for data in train_data:\n",
" train_class_count[data[\"class\"]] += 1\n",
" if data[\"is_broken\"]:\n",
" train_good_defect_count[\"defect\"] += 1\n",
" else:\n",
" train_good_defect_count[\"good\"] += 1\n",
"\n",
"val_class_count = {class_name: 0 for class_name in class_names}\n",
"val_good_defect_count = {\"good\": 0, \"defect\": 0}\n",
"\n",
"for data in val_data:\n",
" val_class_count[data[\"class\"]] += 1\n",
" if data[\"is_broken\"]:\n",
" val_good_defect_count[\"defect\"] += 1\n",
" else:\n",
" val_good_defect_count[\"good\"] += 1\n",
"\n",
"test_class_count = {class_name: 0 for class_name in class_names}\n",
"test_good_defect_count = {\"good\": 0, \"defect\": 0}\n",
"\n",
"for data in test_data:\n",
" test_class_count[data[\"class\"]] += 1\n",
" if data[\"is_broken\"]:\n",
" test_good_defect_count[\"defect\"] += 1\n",
" else:\n",
" test_good_defect_count[\"good\"] += 1\n"
]
},
{
"cell_type": "code",
2025-01-13 04:07:02 +00:00
"execution_count": 17,
2025-01-12 00:43:42 +00:00
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Text(0.5, 0, 'Good/Defect')"
]
},
2025-01-13 04:07:02 +00:00
"execution_count": 17,
2025-01-12 00:43:42 +00:00
"metadata": {},
"output_type": "execute_result"
},
{
"data": {
"image/png": "iVBORw0KGgoAAAANSUhEUgAABkYAAAZZCAYAAAAIy62zAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy88F64QAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOzdebwVBf0//tdlu6wXBGVTFndARQ23q6WkKCru9EmNEv26pIGmqCllLqRhLmkW2maSFVlaWmouaIGl4IKZa7ikYh9Z+qhsLqzz+8MfJy/7ZbvAeT4fj/N4cGbmzLznnDlz583rzExFURRFAAAAAAAAykC9ui4AAAAAAABgXRGMAAAAAAAAZUMwAgAAAAAAlA3BCAAAAAAAUDYEIwAAAAAAQNkQjAAAAAAAAGVDMAIAAAAAAJQNwQgAAAAAAFA2BCMAAAAAAEDZEIwA1KETTzwxXbt2rbPl9+7dO717966z5QMAAKysN954IxUVFRk5cuQana++CKD8CEYAlqKiomKlHmPGjKnrUpdq6tSpOe+889KtW7c0bdo0zZo1S69evXL55Zdn+vTpdV3eStkQ1mHUqFG5/vrr67oMAABY7xxxxBFp2rRpZs2atcxpBgwYkEaNGuWdd95ZKzVsCD3FimwI66AvAjZEDeq6AID10S9+8Ysaz2+99daMHj16ieHdu3dfreX85Cc/ycKFC1drHot78sknc+ihh2b27Nn54he/mF69eiVJnnrqqVx55ZV55JFH8uCDD67RZa5pG8o6jBo1Ks8//3zOPvvsui4FAADWKwMGDMjdd9+dO++8MyeccMIS4z/44IP84Q9/yMEHH5w2bdqs8eVvKD3F8mwo66AvAjZEghGApfjiF79Y4/n48eMzevToJYYv7oMPPkjTpk1XejkNGzZcpfqWZfr06Tn66KNTv379/P3vf0+3bt1qjL/iiivyk5/8ZI0uc03bGNYBAADK3RFHHJEWLVpk1KhRSw1G/vCHP+T999/PgAED1viyN4aeYmNYB4D1mUtpAayi3r17Z8cdd8yECROy7777pmnTpvn617+e5OOD/H79+qVjx46prKzM1ltvnW9961tZsGBBjXksfo+RRdfMveaaa/LjH/84W2+9dSorK7P77rvnySefXGFNP/rRj/K///u/+e53v7vEgXOStGvXLhdddNEyXz937txcfPHF6dWrV1q2bJlmzZrlM5/5TP7yl78sMe1tt92WXr16pUWLFqmqqspOO+2U733ve6Xx8+bNy2WXXZZtt902jRs3Tps2bfLpT386o0ePXuPrcOONN2aHHXZIZWVlOnbsmEGDBi1xWnnXrl1z4oknLjG/xa8nPGbMmFRUVOS3v/1trrjiimyxxRZp3LhxDjjggLz66qs1XnfvvffmzTffLF1arS7vFwMAAOuTJk2a5JhjjsnDDz+cadOmLTF+1KhRadGiRY444oi8++67Oe+887LTTjulefPmqaqqyiGHHJJ//OMfq7RsfZG+CGBFnDECsBreeeedHHLIITnuuOPyxS9+Me3atUuSjBw5Ms2bN8+QIUPSvHnz/PnPf87FF1+cmTNn5uqrr17hfEeNGpVZs2bly1/+cioqKnLVVVflmGOOyb/+9a/lnmXyxz/+MU2aNMnnPve5VVqfmTNn5qc//WmOP/74nHrqqZk1a1Zuvvnm9O3bN0888UR22WWXJMno0aNz/PHH54ADDsh3vvOdJMlLL72URx99NF/96leTJJdeemmGDx+eU045JXvssUdmzpyZp556Kk8//XQOPPDANbYOl156aS677LL06dMnZ5xxRiZOnJibbropTz75ZB599NFVPivnyiuvTL169XLeeedlxowZueqqqzJgwIA8/vjjSZJvfOMbmTFjRv7973/nuuuuS5I0b958lZYFAAAbowEDBuTnP/95fvvb32bw4MGl4e+++24eeOCBHH/88WnSpEleeOGF3HXXXfmf//mfbLnllpk6dWp+9KMfZb/99suLL76Yjh071mq5+iJ9EcCKCEYAVsOUKVPywx/+MF/+8pdrDB81alSaNGlSen766afn9NNPz4033pjLL788lZWVy53vpEmT8sorr2STTTZJkmy//fY58sgj88ADD+Swww5b5uteeumlbLfddmnUqNEqrc8mm2ySN954o8brTz311HTr1i3f//73c/PNNydJ7r333lRVVeWBBx5I/fr1lzqve++9N4ceemh+/OMf16qG2qzDf/7znwwfPjwHHXRQ7rvvvtSr9/GJkN26dcvgwYPzy1/+MieddFKtlr/IRx99lGeeeaZUxyabbJKvfvWref7557PjjjvmwAMPzOabb5733ntvhZdYAwCAcrT//vunQ4cOGTVqVI1g5Pbbb8+8efNKl9Haaaed8vLLL5eO55PkS1/6Urp165abb7453/zmN2u1XH2RvghgRVxKC2A1VFZWLvUA85OhyKxZs/J///d/+cxnPpMPPvgg//znP1c432OPPbYUiiTJZz7zmSTJv/71r+W+bubMmWnRosXKlr+E+vXrlw54Fy5cmHfffTfz58/Pbrvtlqeffro0XatWrfL+++8v9/TvVq1a5YUXXsgrr7xSqxpqsw4PPfRQ5s6dm7PPPrtGE3Xqqaemqqoq9957b62W/UknnXRSjSZkZT8DAADgY/Xr189xxx2XcePG5Y033igNHzVqVNq1a5cDDjggycd91aLj+QULFuSdd95J8+bNs/3229foQ1aWvuhj+iKAZROMAKyGzTfffKm/4HnhhRdy9NFHp2XLlqmqqspmm21W+vXMjBkzVjjfzp0713i+KCR57733lvu6qqqqzJo1a2XLX6qf//zn6dmzZ+n6t5tttlnuvffeGnV/5StfyXbbbZdDDjkkW2yxRf7f//t/uf/++2vMZ9iwYZk+fXq222677LTTTjn//PPz7LPPrnD5tVmHN998M8nHZ9R8UqNGjbLVVluVxq+KVf0MAACA/1p0VsioUaOSJP/+97/z17/+Nccdd1zpLIuFCxfmuuuuy7bbbpvKyspsuumm2WyzzfLss8+uVP+0OH3Rx/RFAMsmGAFYDZ88M2SR6dOnZ7/99ss//vGPDBs2LHfffXdGjx5duubswoULVzjfZZ2GXRTFcl/XrVu3vPzyy5k7d+5KVL+kX/7ylznxxBOz9dZb5+abb87999+f0aNHZ//9969Rd9u2bfPMM8/kj3/8Y4444oj85S9/ySGHHJKBAweWptl3333z2muv5Wc/+1l23HHH/PSnP82nPvWp/PSnP12r67AsFRUVSx2+YMGCpQ5f1c8AAAD4r169eqVbt2759a9/nST59a9/naIoSoFJknz729/OkCFDsu++++aXv/xlHnjggYwePTo77LDDSvVPi9MXLZu+COBjghGANWzMmDF55513MnLkyHz1q1/NYYcdlj59+tS4NNbacvjhh+fDDz/M7373u1V6/R133JGtttoqv//97/OlL30pffv2TZ8+ffLRRx8tMW2jRo1y+OGH58Ybb8xrr72WL3/5y7n11lvz6quvlqZp3bp1TjrppPz617/OW2+9lZ49e+bSSy9dY+vQpUuXJMnEiRNrDJ87d25ef/310vjk4182TZ8+fYl5rM6vp5bVVAAAAP81YMCAPP/883n22WczatSobLvtttl9991L4++444589rOfzc0335zjjjsuBx10UPr06bPU4/eVoS/6mL4IYNkEIwBr2KJf1HzyFzRz587NjTfeuNaXffrpp6dDhw4599xz8/LLLy8xftq0abn88suX+fql1f74449n3LhxNaZ75513ajyvV69eevbsmSSZM2fOUqdp3rx5ttlmm9L4NbEOffr0SaNGjXLDDTfUqPnmm2/OjBkz0q9fv9KwrbfeOuPHj6/xi6t77rknb7311nLrWZ5mzZqt0qn9AABQThadHXLxxRfnmWeeqXG2SPJxH7L4GQi33357/vd//3eVlqcv+pi+CGDZGtR1AQAbm7333jubbLJJBg4cmLPOOisVFRX5xS9+sU5ONd5kk01y55135tBDD80uu+ySL37xi+nVq1eS5Omnn86
"text/plain": [
"<Figure size 2000x2000 with 6 Axes>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# plot the data balance\n",
"\n",
"import matplotlib.pyplot as plt\n",
"\n",
"fig, ax = plt.subplots(3, 2, figsize=(20, 20))\n",
"\n",
"ax[0, 0].bar(train_class_count.keys(), train_class_count.values())\n",
"ax[0, 0].set_title(\"Train Class Count\")\n",
"ax[0, 0].set_xlabel(\"Class\")\n",
"\n",
"ax[0, 1].bar(val_class_count.keys(), val_class_count.values())\n",
"ax[0, 1].set_title(\"Val Class Count\")\n",
"ax[0, 1].set_xlabel(\"Class\")\n",
"\n",
"ax[1, 0].bar(test_class_count.keys(), test_class_count.values())\n",
"ax[1, 0].set_title(\"Test Class Count\")\n",
"ax[1, 0].set_xlabel(\"Class\")\n",
"\n",
"ax[1, 1].bar(train_good_defect_count.keys(), train_good_defect_count.values())\n",
"ax[1, 1].set_title(\"Train Good Defect Count\")\n",
"ax[1, 1].set_xlabel(\"Good/Defect\")\n",
"\n",
"ax[2, 0].bar(val_good_defect_count.keys(), val_good_defect_count.values())\n",
"ax[2, 0].set_title(\"Val Good Defect Count\")\n",
"ax[2, 0].set_xlabel(\"Good/Defect\")\n",
"\n",
"ax[2, 1].bar(test_good_defect_count.keys(), test_good_defect_count.values())\n",
"ax[2, 1].set_title(\"Test Good Defect Count\")\n",
"ax[2, 1].set_xlabel(\"Good/Defect\")\n",
"\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "minigptv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.21"
}
},
"nbformat": 4,
"nbformat_minor": 2
}