Update utils.py
Browse files
utils.py
CHANGED
|
@@ -52,6 +52,7 @@ def create_annotated_video(video_path, df, mse_embeddings, largest_cluster, outp
|
|
| 52 |
if not ret:
|
| 53 |
break
|
| 54 |
# Detect face and draw bounding box
|
|
|
|
| 55 |
mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_face_size=50)
|
| 56 |
boxes, _ = mtcnn.detect(frame)
|
| 57 |
if boxes is not None and len(boxes) > 0:
|
|
|
|
| 52 |
if not ret:
|
| 53 |
break
|
| 54 |
# Detect face and draw bounding box
|
| 55 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 56 |
mtcnn = MTCNN(keep_all=False, device=device, thresholds=[0.9, 0.9, 0.9], min_face_size=50)
|
| 57 |
boxes, _ = mtcnn.detect(frame)
|
| 58 |
if boxes is not None and len(boxes) > 0:
|