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()