Create collect_data.py

collect_data.py

@@ -0,0 +1,113 @@
+# collect_data.py
+import cv2
+import mediapipe as mp
+import numpy as np
+import os
+import time
+from collections import deque
+
+# configuration
+DATA_DIR = "gesture_data"
+SEQUENCE_LENGTH = 30  # number of frames per sample
+EXAMPLES_PER_LABEL = 50
+LABELS = ["air_lock", "swipe_left", "swipe_right", "circle", "hug"]  # update as needed
+
+os.makedirs(DATA_DIR, exist_ok=True)
+
+mp_hands = mp.solutions.hands
+mp_drawing = mp.solutions.drawing_utils
+
+def extract_landmarks(hand_landmarks):
+    # returns 21*3 normalized coords (x,y,z) flattened; if hand missing, return zeros
+    if hand_landmarks is None:
+        return np.zeros(21 * 3, dtype=np.float32)
+    coords = []
+    for lm in hand_landmarks.landmark:
+        coords.extend([lm.x, lm.y, lm.z])
+    return np.array(coords, dtype=np.float32)
+
+def capture_label_sequence(label, cap, hands):
+    seq = deque(maxlen=SEQUENCE_LENGTH)
+    print(f"Prepare to record label: {label}. Press 'r' to start recording one example.")
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            raise RuntimeError("Failed reading webcam")
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        res = hands.process(frame_rgb)
+        left_hand = None
+        right_hand = None
+        # We will just use first detected hand (or zeros) to keep it simple:
+        lm_vec = None
+        if res.multi_hand_landmarks:
+            # choose the first hand
+            lm_vec = extract_landmarks(res.multi_hand_landmarks[0])
+            mp_drawing.draw_landmarks(frame, res.multi_hand_landmarks[0], mp_hands.HAND_CONNECTIONS)
+        else:
+            lm_vec = extract_landmarks(None)
+        cv2.putText(frame, f"Label: {label} | Press 'r' start, 'q' quit", (10,30),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)
+        cv2.imshow("Collect Gestures", frame)
+        k = cv2.waitKey(1) & 0xFF
+        if k == ord('r'):
+            # record one sequence
+            seq.clear()
+            print("Recording...")
+            t0 = time.time()
+            while len(seq) < SEQUENCE_LENGTH:
+                ret, f2 = cap.read()
+                if not ret: break
+                f2_rgb = cv2.cvtColor(f2, cv2.COLOR_BGR2RGB)
+                r = hands.process(f2_rgb)
+                if r.multi_hand_landmarks:
+                    vec = extract_landmarks(r.multi_hand_landmarks[0])
+                    mp_drawing.draw_landmarks(f2, r.multi_hand_landmarks[0], mp_hands.HAND_CONNECTIONS)
+                else:
+                    vec = extract_landmarks(None)
+                seq.append(vec)
+                cv2.putText(f2, f"Recording... {len(seq)}/{SEQUENCE_LENGTH}", (10,30),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,255), 2)
+                cv2.imshow("Collect Gestures", f2)
+                cv2.waitKey(1)
+            t1 = time.time()
+            print(f"Finished recording (took {t1-t0:.2f}s).")
+            if len(seq) == SEQUENCE_LENGTH:
+                arr = np.stack(seq, axis=0)  # (seq_len, features)
+                # save into disk
+                label_dir = os.path.join(DATA_DIR, label)
+                os.makedirs(label_dir, exist_ok=True)
+                idx = len(os.listdir(label_dir))
+                fname = os.path.join(label_dir, f"{idx:04d}.npz")
+                np.savez_compressed(fname, data=arr)
+                print(f"Saved {fname}")
+                return True
+        elif k == ord('q'):
+            return False
+
+def main():
+    cap = cv2.VideoCapture(0)
+    if not cap.isOpened():
+        raise RuntimeError("Cannot open webcam")
+    with mp_hands.Hands(static_image_mode=False,
+                        max_num_hands=1,
+                        min_detection_confidence=0.5,
+                        min_tracking_confidence=0.5) as hands:
+        for label in LABELS:
+            label_dir = os.path.join(DATA_DIR, label)
+            os.makedirs(label_dir, exist_ok=True)
+            cur = len(os.listdir(label_dir))
+            print(f"Label '{label}' currently has {cur} examples. Target: {EXAMPLES_PER_LABEL}")
+            while cur < EXAMPLES_PER_LABEL:
+                ok = capture_label_sequence(label, cap, hands)
+                if not ok:
+                    print("User requested quit.")
+                    cap.release()
+                    cv2.destroyAllWindows()
+                    return
+                cur = len(os.listdir(label_dir))
+                print(f"Now {cur}/{EXAMPLES_PER_LABEL} for label '{label}'")
+    cap.release()
+    cv2.destroyAllWindows()
+
+if __name__ == "__main__":
+    main()