Spaces:
Sleeping
Sleeping
File size: 1,445 Bytes
de5332a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | 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]
|