Introduction

  • なんか突然メガネのセグメンテーションしたくなったのでやってみました。
  • glasses-detectorというまんまのパッケージがあったのでそれを利用

ソースコード

  • OpenCVでWebカメラをキャプチャして、フレーム単位で読み込みます。
  • 読み込んだフレームをglasses_detectorに推論させて、マスクを生成
  • マスクを画像に重ねて、結果表示させる感じです。
import cv2  
import torch  
from glasses_detector import GlassesSegmenter  
import numpy as np  
  
print(f"CUDA:{torch.cuda.is_available()}")  
  
def main():  
    glass_segmenter = GlassesSegmenter(kind="lenses")  
    # Webカメラをキャプチャ  
    cap = cv2.VideoCapture(0)  
  
    while True:  
        # フレームを読み込む  
        ret, frame = cap.read()  
        if not ret:  
            break  
  
        # mask部分推論  
        masks = glass_segmenter.predict(frame,format="mask")  
        # mask描画  
        if cv2.waitKey(1) & 0xFF == ord('m'):  
            results = glass_segmenter.draw_masks(frame,masks,alpha=1,colors=(255,255,255))  
        else:  
            results = glass_segmenter.draw_masks(frame,masks,alpha=0,colors=(255,255,255))  
        # 結果を表示  
        cv2.imshow('Object Detection', np.array(results))  
  
        # 'q'キーで終了  
        if cv2.waitKey(1) & 0xFF == ord('q'):  
            break  
  
    # リソースを解放  
    cap.release()  
    cv2.destroyAllWindows()  
  
  
if __name__ == '__main__':  
    main()
 
  • このソースコードで出力されるのは、こんな感じ。