raab/efishSignalDetector
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

adapted data structure

Browse Source
This commit is contained in:
Till Raab 2023-11-07 15:58:10 +01:00
parent 28554efe2b
commit 9dcf54c139
3 changed files with 93 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

89
datasets.py
View File

@ -14,11 +14,13 @@ from pathlib import Path
from tqdm.auto import tqdm
from PIL import Image
from confic import (CLASSES, RESIZE_TO, DATA_DIR, BATCH_SIZE)
from confic import (CLASSES, RESIZE_TO, DATA_DIR, LABEL_DIR, BATCH_SIZE, IMG_SIZE, IMG_DPI)
from custom_utils import collate_fn
from IPython import embed
from sklearn.model_selection import train_test_split
class InferenceDataset(Dataset):
def __init__(self, dir_path):
@ -75,6 +77,81 @@ class CustomDataset(Dataset):
return len(self.all_images)
class CustomDataset_v2(Dataset):
def __init__(self, limited_idxs=None):
self.images = np.array(sorted(os.listdir(DATA_DIR)))
self.labels = np.array(sorted(os.listdir(LABEL_DIR)))
if hasattr(limited_idxs, '__len__'):
self.images = np.array(sorted(os.listdir(DATA_DIR)))[limited_idxs]
self.labels = np.array(sorted(os.listdir(LABEL_DIR)))[limited_idxs]
self.file_names = np.array([Path(x).with_suffix('') for x in self.images])
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
img = Image.open(Path(DATA_DIR) / self.images[idx])
img_tensor = F.to_tensor(img.convert('RGB'))
annotations = np.loadtxt(Path(LABEL_DIR) / Path(self.images[idx]).with_suffix('.txt'), delimiter=' ')
boxes, labels, area, iscrowd = self.extract_bboxes(annotations)
target = {}
target["boxes"] = boxes
target["labels"] = torch.as_tensor(labels, dtype=torch.int64)
target["area"] = area
target["iscrowd"] = iscrowd
image_id = torch.tensor([idx])
target["image_id"] = image_id
target["image_name"] = self.images[idx]
return img_tensor, target
def extract_bboxes(self, annotations):
if len(annotations.shape) == 1:
annotations = np.array([annotations])
if annotations.shape[1] == 0:
boxes = area = torch.tensor([], dtype=torch.float32)
labels = iscrowd = torch.tensor([], dtype=torch.int64)
return boxes, labels, area, iscrowd
boxes = np.array([[x[1] - x[3] / 2, x[2] - x[4] / 2, x[1] + x[3] / 2, x[2] + x[4] / 2] for x in annotations])
boxes[:, 0] *= IMG_SIZE[0] * IMG_DPI
boxes[:, 2] *= IMG_SIZE[0] * IMG_DPI
boxes[:, 1] *= IMG_SIZE[1] * IMG_DPI
boxes[:, 3] *= IMG_SIZE[1] * IMG_DPI
boxes = torch.from_numpy(boxes).type(torch.float32)
labels = torch.from_numpy(annotations[:, 0]).type(torch.int64)
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
iscrowd = torch.zeros((boxes.shape[0],), dtype=torch.int64)
return boxes, labels, area, iscrowd
def custom_train_test_split():
file_list = sorted(list(Path(LABEL_DIR).rglob('*.txt')))
data_idxs = np.arange(len(file_list))
empty_mask = np.array([os.stat(x).st_size == 0 for x in file_list], dtype=bool)
data_idxs = data_idxs[~empty_mask]
# ToDo: do this witch labels and remove empty shit !!!
np.random.shuffle(data_idxs)
train_idxs = np.sort(data_idxs[int(0.2 * len(data_idxs)):])
test_idxs = np.sort(data_idxs[:int(0.2 * len(data_idxs))])
train_data = CustomDataset_v2(limited_idxs=train_idxs)
test_data = CustomDataset_v2(limited_idxs=test_idxs)
return train_data, test_data
def create_train_or_test_dataset(path, train=True):
if train == True:
pfx='train'
@ -103,6 +180,7 @@ def create_train_loader(train_dataset, num_workers=0):
return train_loader
# ToDo the next two functions are redundant!
def create_valid_loader(valid_dataset, num_workers=0):
valid_loader = DataLoader(
valid_dataset,
@ -125,16 +203,14 @@ def create_inference_loader(inference_dataset, num_workers=0):
if __name__ == '__main__':
# train_data, test_data = create_train_test_dataset(TRAIN_DIR)
train_data = create_train_or_test_dataset(DATA_DIR)
test_data = create_train_or_test_dataset(DATA_DIR, train=False)
train_data, test_data = custom_train_test_split()
train_loader = create_train_loader(train_data)
test_loader = create_valid_loader(test_data)
for samples, targets in test_loader:
for samples, targets in train_loader:
for s, t in zip(samples, targets):
fig, ax = plt.subplots()
ax.imshow(s.permute(1, 2, 0), aspect='auto')
for (x0, y0, x1, y1), l in zip(t['boxes'], t['labels']):
@ -145,4 +221,5 @@ if __name__ == '__main__':
(y1 - y0),
fill=False, color="white", linewidth=2, zorder=10)
)
ax.set_title(t['image_name'])
plt.show()

4
inference.py
View File

@ -8,7 +8,7 @@ import argparse
from model import create_model
from confic import NUM_CLASSES, DEVICE, CLASSES, OUTDIR, DATA_DIR, INFERENCE_OUTDIR, IMG_DPI, IMG_SIZE
from datasets import InferenceDataset, create_inference_loader
from datasets import InferenceDataset, create_valid_loader
from IPython import embed
from pathlib import Path
@ -66,7 +66,7 @@ def main(args):
model.to(DEVICE).eval()
inference_data = InferenceDataset(args.folder)
inference_loader = create_inference_loader(inference_data)
inference_loader = create_valid_loader(inference_data)
dataset_name = Path(args.folder).name

11
train.py
View File

@ -1,6 +1,6 @@
from confic import (DEVICE, NUM_CLASSES, NUM_EPOCHS, OUTDIR, NUM_WORKERS, DATA_DIR, IMG_SIZE, IMG_DPI, INFERENCE_OUTDIR)
from model import create_model
from datasets import create_train_loader, create_valid_loader, create_train_or_test_dataset
from datasets import create_train_loader, create_valid_loader, custom_train_test_split
from custom_utils import Averager, SaveBestModel, save_model, save_loss_plot
from tqdm.auto import tqdm
@ -121,8 +121,13 @@ def plot_validation(img_tensor, img_name, output, target, detection_threshold):
# plt.show()
if __name__ == '__main__':
train_data = create_train_or_test_dataset(DATA_DIR)
test_data = create_train_or_test_dataset(DATA_DIR, train=False)
# train_data = create_train_or_test_dataset(DATA_DIR)
# test_data = create_train_or_test_dataset(DATA_DIR, train=False)
#
# train_loader = create_train_loader(train_data)
# test_loader = create_valid_loader(test_data)
train_data, test_data = custom_train_test_split()
train_loader = create_train_loader(train_data)
test_loader = create_valid_loader(test_data)