Dumitrana-Mihnea
51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
import os
|
|
import mediapipe as mp
|
|
import cv2
|
|
import pickle
|
|
from sklearn.ensemble import RandomForestClassifier
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.metrics import accuracy_score
|
|
import numpy as np
|
|
|
|
mp_hands = mp.solutions.hands
|
|
hands = mp_hands.Hands(static_image_mode=True, min_detection_confidence=0.5)
|
|
data_dir = "./data"
|
|
data = []
|
|
labels = []
|
|
|
|
for dir_ in os.listdir(data_dir):
|
|
print(dir_)
|
|
for img_path in os.listdir(os.path.join(data_dir, dir_)):
|
|
data_aux = []
|
|
img = cv2.imread(os.path.join(data_dir, dir_, img_path))
|
|
|
|
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
|
|
data_aux.append(x)
|
|
data_aux.append(y)
|
|
# print(f"Poza {img_path} cu ", end=' ')
|
|
# print(x, y)
|
|
data.append(data_aux)
|
|
labels.append(dir_)
|
|
|
|
|
|
data = np.asarray(data)
|
|
labels = np.asarray(labels)
|
|
x_train, x_test, y_train, y_test = train_test_split(data, labels, test_size=0.2, shuffle=True, stratify=labels)
|
|
model = RandomForestClassifier()
|
|
model.fit(x_train, y_train)
|
|
y_predict = model.predict(x_test)
|
|
score = accuracy_score(y_predict, y_test)
|
|
|
|
print(f"{score * 100}% classified correctly")
|
|
f = open("model.p", "wb")
|
|
pickle.dump({"model": model}, f)
|
|
f.close()
|
|
|
|
|