From 313d41ea737b6b7f5a8d8f2b56b4b1cce1440bda Mon Sep 17 00:00:00 2001
From: BigMaster216 <jameskennedy415216@gmail.com>
Date: Tue, 2 Jul 2024 14:21:24 -0500
Subject: [PATCH] Add the test camera script and optimize the create dataset.

---
 collect_imgs.py         |  2 +-
 create_dataset.py       | 76 ++++++++++++++++++++++++++---------------
 inference_classifier.py |  2 +-
 test_camera.py          | 12 +++++++
 4 files changed, 62 insertions(+), 30 deletions(-)
 create mode 100644 test_camera.py

diff --git a/collect_imgs.py b/collect_imgs.py
index bb987be93..43b605be8 100644
--- a/collect_imgs.py
+++ b/collect_imgs.py
@@ -10,7 +10,7 @@
 number_of_classes = 3
 dataset_size = 100
 
-cap = cv2.VideoCapture(2)
+cap = cv2.VideoCapture(0)
 for j in range(number_of_classes):
     if not os.path.exists(os.path.join(DATA_DIR, str(j))):
         os.makedirs(os.path.join(DATA_DIR, str(j)))
diff --git a/create_dataset.py b/create_dataset.py
index cf3aebd9a..54a6eb8a2 100644
--- a/create_dataset.py
+++ b/create_dataset.py
@@ -1,50 +1,70 @@
 import os
 import pickle
-
-import mediapipe as mp
 import cv2
+import mediapipe as mp
 import matplotlib.pyplot as plt
 
-
+# Initialize MediaPipe Hands
 mp_hands = mp.solutions.hands
 mp_drawing = mp.solutions.drawing_utils
-mp_drawing_styles = mp.solutions.drawing_styles
 
 hands = mp_hands.Hands(static_image_mode=True, min_detection_confidence=0.3)
 
+# Define data directory
 DATA_DIR = './data'
 
+# Initialize lists for data and labels
 data = []
 labels = []
-for dir_ in os.listdir(DATA_DIR):
-    for img_path in os.listdir(os.path.join(DATA_DIR, dir_)):
-        data_aux = []
 
-        x_ = []
-        y_ = []
+# Function to process images
+def process_image(img_path):
+    data_aux = []
+    x_ = []
+    y_ = []
 
-        img = cv2.imread(os.path.join(DATA_DIR, dir_, img_path))
-        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = cv2.imread(img_path)
+    if img is None:
+        return None, None
+    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
-        results = hands.process(img_rgb)
-        if results.multi_hand_landmarks:
-            for hand_landmarks in results.multi_hand_landmarks:
-                for i in range(len(hand_landmarks.landmark)):
-                    x = hand_landmarks.landmark[i].x
-                    y = hand_landmarks.landmark[i].y
+    results = hands.process(img_rgb)
+    if results.multi_hand_landmarks:
+        for hand_landmarks in results.multi_hand_landmarks:
+            for i in range(len(hand_landmarks.landmark)):
+                x = hand_landmarks.landmark[i].x
+                y = hand_landmarks.landmark[i].y
+                x_.append(x)
+                y_.append(y)
 
-                    x_.append(x)
-                    y_.append(y)
+            for i in range(len(hand_landmarks.landmark)):
+                x = hand_landmarks.landmark[i].x
+                y = hand_landmarks.landmark[i].y
+                data_aux.append(x - min(x_))
+                data_aux.append(y - min(y_))
 
-                for i in range(len(hand_landmarks.landmark)):
-                    x = hand_landmarks.landmark[i].x
-                    y = hand_landmarks.landmark[i].y
-                    data_aux.append(x - min(x_))
-                    data_aux.append(y - min(y_))
+        return data_aux, True
+    return None, False
 
-            data.append(data_aux)
+# Process each image in the data directory
+for dir_ in os.listdir(DATA_DIR):
+    class_dir = os.path.join(DATA_DIR, dir_)
+    if not os.path.isdir(class_dir):
+        continue
+    
+    for img_name in os.listdir(class_dir):
+        img_path = os.path.join(class_dir, img_name)
+        processed_data, success = process_image(img_path)
+
+        if success:
+            data.append(processed_data)
             labels.append(dir_)
 
-f = open('data.pickle', 'wb')
-pickle.dump({'data': data, 'labels': labels}, f)
-f.close()
+# Save data to a pickle file
+with open('data.pickle', 'wb') as f:
+    pickle.dump({'data': data, 'labels': labels}, f)
+
+# Release resources
+hands.close()
+
+print("Data processing completed and saved to 'data.pickle'.")
diff --git a/inference_classifier.py b/inference_classifier.py
index ce322366c..e90216b8b 100644
--- a/inference_classifier.py
+++ b/inference_classifier.py
@@ -7,7 +7,7 @@
 model_dict = pickle.load(open('./model.p', 'rb'))
 model = model_dict['model']
 
-cap = cv2.VideoCapture(2)
+cap = cv2.VideoCapture(0)
 
 mp_hands = mp.solutions.hands
 mp_drawing = mp.solutions.drawing_utils
diff --git a/test_camera.py b/test_camera.py
new file mode 100644
index 000000000..d3aeeb7cb
--- /dev/null
+++ b/test_camera.py
@@ -0,0 +1,12 @@
+import cv2
+
+def test_cameras():
+    for i in range(5):  # Test indices 0 to 4
+        cap = cv2.VideoCapture(i)
+        if cap.isOpened():
+            print(f"Camera {i} is available.")
+            cap.release()
+        else:
+            print(f"Camera {i} is not available.")
+
+test_cameras()