From 9dcf54c1390f548df063060ab4b5e04cf9668497 Mon Sep 17 00:00:00 2001
From: Till Raab <till.raab@uni-tuebingen.de>
Date: Tue, 7 Nov 2023 15:58:10 +0100
Subject: [PATCH] adapted data structure
---
datasets.py | 89 ++++++++++++++++++++++++++++++++++++++++++++++++----
inference.py | 4 +--
train.py | 11 +++++--
3 files changed, 93 insertions(+), 11 deletions(-)
diff --git a/datasets.py b/datasets.py
index 284efe8..ad5effd 100644
--- a/datasets.py
+++ b/datasets.py
@@ -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()
\ No newline at end of file
diff --git a/inference.py b/inference.py
index a09f0d2..326398c 100644
--- a/inference.py
+++ b/inference.py
@@ -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
diff --git a/train.py b/train.py
index 56b7389..38b93a8 100644
--- a/train.py
+++ b/train.py
@@ -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)