# %%
!pip install opencv-python

# %%
import pandas as pd
train = pd.read_csv("train.csv")

# %%
import cv2
import numpy as np
from sklearn.ensemble import RandomForestRegressor

def extract(img):
    # img = cv2.GaussianBlur(img, (3,3), 0)
    return [
        np.mean(img),
        np.std(img),
        cv2.Laplacian(img, cv2.CV_64F).var()
        ]

X = []
y_count = []
y_age = []

for i in range(700):
    img = cv2.imread(f"images/{i:04d}.png", cv2.IMREAD_GRAYSCALE)
    X.append(extract(img))
    y_count.append(train.loc[i, "count"])
    y_age.append(train.loc[i, "age"])

X = np.array(X)

model_count = RandomForestRegressor()
model_age = RandomForestRegressor()

model_count.fit(X, y_count)
model_age.fit(X, y_age)

# %%
import cv2
import numpy as np

def extract(img):
    # img = cv2.GaussianBlur(img, (7,7), 0)
    return [
        np.mean(img),
        np.std(img),
        cv2.Laplacian(img, cv2.CV_64F).var()
    ]

rows = []

for i in range(700, 1000):
    img = cv2.imread(f"images/{i:04d}.png", cv2.IMREAD_GRAYSCALE)
    feats = np.array(extract(img)).reshape(1, -1)

    pred_count = model_count.predict(feats)[0]
    pred_age = model_age.predict(feats)[0]
    pred_count = int(np.clip(round(pred_count), 3, 8))
    pred_age = float(np.clip(pred_age, 1.0, 5.0))

    rows.append([1, i, pred_count])
    rows.append([2, i, pred_age])

submission = pd.DataFrame(rows, columns=["subtaskID", "datapointID", "answer"])
submission.to_csv("submission.csv", index=False)


