Spaces:
Sleeping
Sleeping
| from sklearn.svm import LinearSVC | |
| import numpy as np | |
| import cv2 as cv | |
| import os | |
| class Model: | |
| def __init__(self): | |
| self.model = LinearSVC() | |
| self.trained = False | |
| def train_model(self, counters): | |
| images = [] | |
| labels = [] | |
| # Class 1 | |
| for i in range(1, counters[0]): | |
| path = f'1/frame{i}.jpg' | |
| if os.path.exists(path): | |
| img = cv.imread(path, cv.IMREAD_GRAYSCALE) | |
| img = cv.resize(img, (150, 150)) | |
| images.append(img.flatten()) | |
| labels.append(1) | |
| # Class 2 | |
| for i in range(1, counters[1]): | |
| path = f'2/frame{i}.jpg' | |
| if os.path.exists(path): | |
| img = cv.imread(path, cv.IMREAD_GRAYSCALE) | |
| img = cv.resize(img, (150, 150)) | |
| images.append(img.flatten()) | |
| labels.append(2) | |
| if len(images) == 0: | |
| print("No training data found!") | |
| return | |
| X = np.array(images) | |
| y = np.array(labels) | |
| self.model.fit(X, y) | |
| self.trained = True | |
| print("Model trained successfully!") | |
| def predict(self, frame): | |
| if not self.trained: | |
| return None | |
| gray = cv.cvtColor(frame, cv.COLOR_RGB2GRAY) | |
| gray = cv.resize(gray, (150, 150)) | |
| gray = gray.flatten() | |
| prediction = self.model.predict([gray]) | |
| return prediction[0] | |