File size: 939 Bytes
e2af51e |
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 |
import joblib
from config import MAIN_MODEL
import pandas as pd
class ViolencePredictor:
def __init__(self):
self.model = joblib.load(MAIN_MODEL)
def _preprocess_data_pdict(self, data: pd.DataFrame) -> pd.DataFrame:
cols_to_drop = [
"video_name",
"frame_index",
"timestamp",
"frame_width",
"frame_height",
"person1_id",
"person2_id",
"person1_idx",
"person2_idx",
]
data = data.drop(columns=cols_to_drop)
return data
def predict(self, data):
data = self._preprocess_data_pdict(data)
y_pred = self.model.predict(data)
print(y_pred)
return y_pred
if __name__ == "__main__":
import pandas as pd
data = pd.read_csv("data/fight_train.csv")
data = data[0:20]
print("dataloaded")
VP = ViolencePredictor()
VP.predict(data)
|