YOLOv5 label conf 수정하기

YOLOv5를 사용하여 모델 학습 중 output의 label 수정이 필요한 경우에 대한 해결법입니다.

 

우선 YOLOv5를 학습시키고 난 후에,

학습된 모델을 사용하여 이미지 결과를 확인하는 단계입니다.

 

yolov5 폴더에 내장된 detect.py 를 사용하여 이미지를 예측하기 위해 아래와 같은 코드를 작성합니다.

!python detect.py --weights "/content/yolov5/runs/train/tree_results/weights/best.pt" --source "/content/test.jpg"

예측할 때는 모델의 경로(–weights) 즉, 학습된 모델의 가중치 파일이 저장된 경로를 입력합니다.

가중치 파일이 저장된 위치는 조금씩 다를 수 있지만, 보통 runs/train 폴더 안에 best.pt 파일이 저장된 경로를 수정해서 적어주시면 됩니다.

다음으로는 예측할 이미지 경로(–source)를 지정해주면 됩니다. 

 

해당 코드를 실행하기 전에 detect.py 파일을 열어 159~171 line에 코드를 부분 수정하면 됩니다.

 for *xyxy, conf, cls in reversed(det):
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or save_crop or view_img:  # Add bbox to image
                        c = int(cls)  # integer class
                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                        annotator.box_label(xyxy, label, color=colors(c, True))
                        if save_crop:
                            save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

 

label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')

위 코드 중에 {conf:.2f} 인 부분을 수정하면 됩니다. 현재는 각 class의 confidence 값을 소수점 두자리로 출력하고 있기 때문에 아래와 같이 0.77 형태로 label이 작성됩니다.

저는 YOLOX와 결과를 비교하기 위해 YOLOX의 label 형태와 같도록 수정했지만, 이를 응용하여 원하는 label 로 수정해서 출력할 수 있습니다.

{names[c]} {conf:.2f} ↠ f'{names[c]}:{conf*100:.1f}%' 

다음과 같이 코드를 수정해 normal:77.1% 로 label의 형태를 변경했습니다. 기존에 소수점인 conf에 100을 곱한 후에 .1f 를 사용해 소수점 1째자리까지 나타냈습니다. 변경된 전체 코드는 다음과 같습니다.

 for *xyxy, conf, cls in reversed(det):
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or save_crop or view_img:  # Add bbox to image
                        c = int(cls)  # integer class
                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]}:{conf*100:.1f}%')
                        annotator.box_label(xyxy, label, color=colors(c, True))
                        if save_crop:
                            save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

 

                              수정 전                                                                     수정 후

 

다음과 같이 수정된 결과를 확인할 수 있습니다. 알려드린 방법으로 label 형태 뿐만 아니라 색상, 크기 또한 수정할 수 있습니다.

 

어려운 코드 수정은 아니지만, label 수정할려고 찾아보니 정리된 블로그가 없어서 작성합니다.

도움되시길 바랍니다! 🌝