-
Notifications
You must be signed in to change notification settings - Fork 0
/
preprocessing.py
40 lines (28 loc) · 1.14 KB
/
preprocessing.py
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
import pandas as pd
import cv2
from sklearn.model_selection import train_test_split
def load_csv(path: str = "dataset/anno_train.csv", nrows: int = None):
"""Read csv file"""
data = pd.read_csv(path, nrows=nrows)
return data
def __preprocess_image(image_path):
image = cv2.imread(f"car_ims/cars_train/{image_path}")
image = cv2.resize(image, (224, 224))
image = image / 255.0 # Normalize pixel values
return image
def prepare(data):
"""Prepare the data to train\n
:return data"""
data['image'] = data['relative_im_path'].apply(__preprocess_image)
# print(len(data.values[:, -1]))
img_width, img_height, num_channels = 224, 224, 4
data = data.reshape(data.shape[0], img_width, img_height, num_channels)
return data
def extract_target(data):
"""Extract targets from data"""
targets = data[['bbox_x1', 'bbox_y1', 'bbox_x2', 'bbox_y2']].values
return targets
def split_train_test(data, targets):
train_data, test_data, train_targets, test_targets = train_test_split(
data['image'].values, targets, test_size=0.2)
return train_data, test_data, train_targets, test_targets