Преглед
Като човек, който е работил в инженерната индустрия повече от 20 години, ръководейки главно екип, фокусиран върху разработката на рамка и автоматизацията, винаги търся начини да използвам нови технологии за подобряване на качеството на продукта. Една област, в която машинното обучение може да бъде особено полезно, е проверката на качеството по време на производството. Например, откриването на модули с дефекти, като драскотини по повърхността или липсващи винтове по време на сглобяването, може да се извърши с помощта на техники за откриване на обекти. Въпреки че има налични предварително обучени модели за откриване на обекти, фината настройка на тези модели с потребителски данни може да доведе до значително по-добра производителност. В тази статия ще предоставя пълно ръководство стъпка по стъпка за фина настройка на модел R-CNN ResNet-50 с персонализирани данни, така че да можете да започнете да използвате откриването на обекти, за да подобрите собствените си процеси за проверка на качеството.
Какво е фина настройка на модел?
С прости думи, това се отнася до процеса на вземане на предварително обучен модел за машинно обучение, като R-CNN ResNet-50, и допълнителното му обучение върху нов набор от данни или задача, за да постигнете вашата конкретна цел. В моя случай искам ResNet-50 да открива винтове, конектори и подложки вместо котки, кучета и лисици. Следното изображение илюстрира процеса на фина настройка на предварително обучен модел на машинно обучение върху нов набор от данни.
Въпреки че има много статии, достъпни за обществеността, които обсъждат тази тема, не можах да намеря изчерпателно ръководство, което предоставя инструкции стъпка по стъпка по време на моето проучване. Ето защо написах тази статия. Бих искал да изразя благодарността си към автора на статията „Как да обучите детектор на обекти с вашия собствен COCO набор от данни в PyTorch (Общи обекти в контекстен формат)““, тъй като ми предостави безценни насоки и идеи, които помогнаха да довърша тази статия.
Подготовка на данни и анотация
За да позволите на R-CNN ResNet-50 да разпознае вашите обекти, напр. винтове или подложка, трябва да направите следните 2 стъпки:
- За да разрешите на R-CNN ResNet-50 да разпознава вашите обекти, като винтове или подложки, трябва да следвате тези две стъпки:
- Подгответе набор от изображения: Според мен 100 изображения са достатъчни, за да преквалифицирате модела.
Анотирайте изображенията с обектите, които искате да откриете: Има много безплатни инструменти за анотиране, като например https://www .makesense.ai/. След като анотацията е готова, можете да експортирате Coco JSON файла.
Анотирането едно по едно на всяко от 100-те различни изображения може да бъде досаден процес. За да опростя това, вместо това използвах подобни изображения. По този начин бих могъл да поясня само едно изображение и да използвам Coco JSON с едно изображение, за да генерирам програмно окончателния Coco JSON, който ще покрие всички изображения.
Фина настройка на RCNN модел с дадения от вас Coco JSON и изображения
Създайте собствен клас от набор от данни
import os
import torch
import torch.utils.data
import torchvision
from PIL import Image
from pycocotools.coco import COCO
class myOwnDataset(torch.utils.data.Dataset):
def __init__(self, root, annotation, transforms=None):
self.root = root
self.transforms = transforms
self.coco = COCO(annotation)
self.ids = list(sorted(self.coco.imgs.keys()))
def __getitem__(self, index):
# Own coco file
coco = self.coco
# Image ID
img_id = self.ids[index]
# List: get annotation id from coco
ann_ids = coco.getAnnIds(imgIds=img_id)
# Dictionary: target coco_annotation file for an image
coco_annotation = coco.loadAnns(ann_ids)
# path for input image
path = coco.loadImgs(img_id)[0]['file_name']
# open the input image
img = Image.open(os.path.join(self.root, path))
# number of objects in the image
num_objs = len(coco_annotation)
# Bounding boxes for objects
# In coco format, bbox = [xmin, ymin, width, height]
# In pytorch, the input should be [xmin, ymin, xmax, ymax]
boxes = []
labels = []
for i in range(num_objs):
xmin = coco_annotation[i]['bbox'][0]
ymin = coco_annotation[i]['bbox'][1]
xmax = xmin + coco_annotation[i]['bbox'][2]
ymax = ymin + coco_annotation[i]['bbox'][3]
boxes.append([xmin, ymin, xmax, ymax])
labels.append(coco_annotation[i]['category_id'])
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor(labels, dtype=torch.int64)
# Labels (In my case, I only one class: target class or background)
#labels = torch.ones((num_objs,), dtype=torch.int64)
# Tensorise img_id
img_id = torch.tensor([img_id])
# Size of bbox (Rectangular)
areas = []
for i in range(num_objs):
areas.append(coco_annotation[i]['area'])
areas = torch.as_tensor(areas, dtype=torch.float32)
# Iscrowd
iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
# Annotation is in dictionary format
my_annotation = {}
my_annotation["boxes"] = boxes
my_annotation["labels"] = labels
my_annotation["image_id"] = img_id
my_annotation["area"] = areas
my_annotation["iscrowd"] = iscrowd
if self.transforms is not None:
img = self.transforms(img)
return img, my_annotation
def __len__(self):
return len(self.ids)
Създайте свързани методи
# added ToTensor
def get_transform():
custom_transforms = []
custom_transforms.append(torchvision.transforms.ToTensor())
return torchvision.transforms.Compose(custom_transforms)
# collate_fn needs for batch
def collate_fn(batch):
return tuple(zip(*batch))
Заредете персонализирани изображения и Coco JSON
train_data_dir = <your image location> train_coco = <your coco JSON location>
Моля, обърнете внимание, че файлът train_coco е един файл. Ако сте анотирали множество изображения и сте генерирали множество Coco JSON файлове, можете просто да ги комбинирате, като промените например ID на изображението и ID на анотацията. За повече подробности, моля, вижте урока „Създаване на COCO анотации от нулата“.
Задайте размера на партидата и създайте устройство за зареждане на данни
# Batch size
train_batch_size = 1
# own DataLoader
data_loader = torch.utils.data.DataLoader(my_dataset,
batch_size=train_batch_size,
shuffle=True,
num_workers=4,
collate_fn=collate_fn)
Настройте устройството да използва GPU, ако е приложимо
# select device (whether GPU or CPU)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
За да проверя дали данните, включително изображения и анотации от Coco JSON, са заредени правилно, направих оператор за печат по-долу, за да мога да проверя визуално данните преди обучение.
counter = 1
for imgs, annotations in data_loader:
imgs = list(img.to(device) for img in imgs)
annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
print ("counter", counter)
print(annotations)
counter += 1
Импортирайте и заредете RCNN модел, за да се подготвите за фина настройка
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor def get_model_instance_segmentation(num_classes): # load an instance segmentation model pre-trained pre-trained on COCO model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None) # get number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
Дефинирайте num_classes и създайте моделен обект. В моята анотация имам осем категории в Coco JSON, така че крайният брой класове трябва да бъде # категории + един. В моя случай е девет.
num_classes = 9
num_epochs = 10
model = get_model_instance_segmentation(num_classes)
Преместете модела на правилното устройство и започнете обучението
# move model to the right device
model.to(device)
# parameters
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
len_dataloader = len(data_loader)
# Train model
for epoch in range(num_epochs):
model.train()
i = 0
for imgs, annotations in data_loader:
i += 1
imgs = list(img.to(device) for img in imgs)
annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
loss_dict = model(imgs, annotations)
losses = sum(loss for loss in loss_dict.values())
optimizer.zero_grad()
losses.backward()
optimizer.step()
print(f'Iteration: {i}/{len_dataloader}, Loss: {losses}')
print("Finish training")
Запазете обучения модел
torch.save(model.state_dict(), "faster_rcnn_resnet50_fpn_custom.pth")
Заредете вашия персонализиран модел и приложете
След като вече имате персонализиран модел RCNN_ResNet50, който може да открива вашите обекти, следващата стъпка е да заредите персонализирания модел и да тествате откриването на обекти. Важно е да се отбележи, че тъй като обучихме модела с осем категории, броят на класовете трябва да бъде зададен като девет при тестване на модела.
num_classes = 9
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None, num_classes=num_classes)
model.load_state_dict(torch.load('faster_rcnn_resnet50_fpn_custom.pth')
Преместете модела на GPU, ако има такъв
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
Определете пътя на входното тестово изображение и заредете изображението
img_path = "my_test_image.jpeg"
img = cv2.imread(img_path)
# Convert the image from BGR to RGB
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Convert the image to a PyTorch tensor
img = torch.from_numpy(img).permute(2, 0, 1).float().div(255).unsqueeze(0)
model.eval()
# Move the image to GPU if available
img = img.to(device)
След като приключи, можем да извършим откриването на обект
# Perform object detection
with torch.no_grad():
output = model(img)
След това можем да начертаем изображението с откритите обекти и да изхвърлим ограничителните полета в JSON файл, наречен inference_annotation_json.
# Get the bounding boxes, labels, and scores of the objects in the image
boxes = output[0]['boxes'].cpu().numpy()
labels = output[0]['labels'].cpu().numpy()
scores = output[0]['scores'].cpu().numpy()
# Plot the image with the detected objects
plt.imshow(img.permute(0, 2, 3, 1).squeeze().cpu().numpy())
annotated_dict = {"img": img_path, "inference_annotations": []}
#get category mapping
category_mapping_dict = creating_mapping()
for i in range(boxes.shape[0]):
if scores[i] > 0.5:
box = boxes[i, :]
label = labels[i]
category = category_mapping_dict[label]
score = scores[i]
plt.gca().add_patch(plt.Rectangle((box[0], box[1]), box[2]-box[0], box[3]-box[1], fill=False, color='red'))
plt.text(box[0], box[1], f"{category}({label}): {score:.2f}", fontsize=8, bbox=dict(facecolor='yellow', alpha=0.5))
annotated_dict["inference_annotations"].append({"bbox": box.tolist(), "category": str(category), "score": float(score)})
plt.axis('off')
plt.savefig('output_annotation.jpg')
plt.show()
inference_annotation_json = json.dumps(annotated_dict)
print(inference_annotation_json)
Заключение
Машинното обучение става все по-разпространено в различни области. Надявам се, че тази статия е била полезна в преследването на вашите интереси и поддържането на страстта ви в тази област. Въпреки че не се смятам за експерт, ще продължа да преследвам страстта си към AI/ML и ще я прилагам по практически начини. Не забравяйте, че ключът към успеха в машинното обучение е да продължите да учите и изследвате.
