[코드] yolov5 구현

간단하게 yolov5 코드 구현하기

프레임워크 : pytorch

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

pre-trained 모델을 정의한다.

# Load a pre-trained model
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

클래스는 2개로 설정  classifier 는 fastRCNN

# Replace the classifier with a new one
num_classes = 2  # Number of classes in your dataset
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

evaluation 모드로 설정해 사용하지 않는 layer는 알아서 off하도록한다.

# Put the model in evaluation mode
model.eval()

예시 출력

# Perform a forward pass on an example image
x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
outputs = model(x)

# Outputs is a dictionary containing the model's predictions
print(outputs)