diff --git a/common/research_common/constants.py b/common/research_common/constants.py
index 3b48b728374638f828ee4d624dd2c32e481d5e25..3bfc5cb835e74ea238d7f502a5f3a482cd46432c 100644
--- a/common/research_common/constants.py
+++ b/common/research_common/constants.py
@@ -1,7 +1,11 @@
from pathlib import Path
+DSET_DIR = Path(__file__).parent.parent.parent / "dataset"
-DSET_DIR = Path(__file__).parent.parent.parent / 'dataset'
+TWITCH_DSET_DIR: Path = DSET_DIR / "twitch"
+ROCO_DSET_DIR: Path = DSET_DIR / "dji_roco"
+TENSORFLOW_RECORDS_DIR: Path = DSET_DIR / "tf_records"
-TWITCH_DSET: Path = DSET_DIR / 'twitch'
-ROCO_DSET: Path = DSET_DIR / 'dji_roco'
+TWITCH_DSET_DIR.mkdir(parents=True, exist_ok=True)
+ROCO_DSET_DIR.mkdir(parents=True, exist_ok=True)
+TENSORFLOW_RECORDS_DIR.mkdir(parents=True, exist_ok=True)
diff --git a/common/research_common/dataset/twitch/aerial_view_detector.py b/common/research_common/dataset/twitch/aerial_view_detector.py
index de15b43542a879b1dc36f57acb8dc2dcf75a7d6e..17ce4d75c18c42655ae9cf8fbeef7013c912db96 100644
--- a/common/research_common/dataset/twitch/aerial_view_detector.py
+++ b/common/research_common/dataset/twitch/aerial_view_detector.py
@@ -2,11 +2,11 @@ from pathlib import Path
from skimage import io
-from research_common.constants import TWITCH_DSET
+from research_common.constants import TWITCH_DSET_DIR
from research_common.dataset.twitch.mask_detector import MaskDetector
aerial_view_detector = MaskDetector(
- Path(__file__).parent / 'mask_aerial.jpg',
+ Path(__file__).parent / "mask_aerial.jpg",
[
(527, 528, 292, 297, 20),
(527, 531, 303, 303, 20),
@@ -14,10 +14,10 @@ aerial_view_detector = MaskDetector(
(536, 541, 302, 303, 20),
(543, 544, 292, 297, 20),
(535, 535, 292, 297, 20),
- ]
+ ],
)
-if __name__ == '__main__':
- for file_path in sorted((TWITCH_DSET / 'robots-views').glob('*.jpg')):
+if __name__ == "__main__":
+ for file_path in sorted((TWITCH_DSET_DIR / "robots-views").glob("*.jpg")):
if aerial_view_detector.is_matching(io.imread(str(file_path))):
print(file_path.name)
diff --git a/common/research_common/dataset/twitch/robots_views_extractor.py b/common/research_common/dataset/twitch/robots_views_extractor.py
index 2f24a49bc9126896763e7351d9e21fe69bad6e39..a702a0d8279ccdd3527276219f569cfb4ecc400a 100644
--- a/common/research_common/dataset/twitch/robots_views_extractor.py
+++ b/common/research_common/dataset/twitch/robots_views_extractor.py
@@ -5,11 +5,11 @@ import ffmpeg
import numpy as np
from tqdm import tqdm
-from research_common.constants import TWITCH_DSET
+from research_common.constants import TWITCH_DSET_DIR
from research_common.dataset.twitch.mask_detector import is_image_from_robot_view
from research_common.dataset.twitch.video_frame_generator import VideoFrameGenerator
-RES_DIR: Path = TWITCH_DSET / 'robots-views'
+RES_DIR: Path = TWITCH_DSET_DIR / "robots-views"
RES_DIR.mkdir(parents=True, exist_ok=True)
@@ -19,25 +19,28 @@ class RobotsViewExtractor:
def __init__(self, video_name: str):
self.video_name: str = video_name
- self.video_path = TWITCH_DSET / 'videos' / f'{video_name}.mp4'
+ self.video_path = TWITCH_DSET_DIR / "videos" / f"{video_name}.mp4"
self.frame_generator: VideoFrameGenerator = VideoFrameGenerator(self.video_path, self.FPS)
+ self.count = 0
def run(self):
for i, frame in tqdm(
enumerate(self.frame_generator.generate()),
total=self._get_number_of_frames(),
- desc=f'Extracting robots views from video {self.video_name}.mp4',
- unit='frames',
+ desc=f"Extracting robots views from video {self.video_name}.mp4",
+ unit="frames",
ncols=200,
):
self._process_frame(frame, i)
+ print(f"Detected {self.count} robots views")
def _process_frame(self, frame: np.ndarray, frame_number: int):
if is_image_from_robot_view(frame):
self._save_frame(frame, frame_number)
+ self.count += 1
def _save_frame(self, frame: np.ndarray, frame_number: int):
cv2.imwrite(f"{RES_DIR}/{self.video_name}-frame-{frame_number + 1:06}.jpg", frame)
def _get_number_of_frames(self):
- return int(ffmpeg.probe(str(self.video_path))['format']['duration'].split('.')[0]) * self.FPS
+ return int(ffmpeg.probe(str(self.video_path))["format"]["duration"].split(".")[0]) * self.FPS
diff --git a/common/research_common/scripts/move_aerial_views.py b/common/research_common/scripts/move_aerial_views.py
index 0ed6b86d224dd3eb16e6b0003229aac3f7c87034..64b40198892b7673211f3c1895bd816e32acde70 100644
--- a/common/research_common/scripts/move_aerial_views.py
+++ b/common/research_common/scripts/move_aerial_views.py
@@ -3,20 +3,20 @@ from shutil import move
from skimage import io
from tqdm import tqdm
-from research_common.constants import TWITCH_DSET
+from research_common.constants import TWITCH_DSET_DIR
from research_common.dataset.twitch.aerial_view_detector import aerial_view_detector
-ROBOTS_VIEWS_DIR = TWITCH_DSET / 'robots-views'
-AERIAL_VIEWS_DIR = TWITCH_DSET / 'aerial-views'
+ROBOTS_VIEWS_DIR = TWITCH_DSET_DIR / "robots-views"
+AERIAL_VIEWS_DIR = TWITCH_DSET_DIR / "aerial-views"
AERIAL_VIEWS_DIR.mkdir(parents=True, exist_ok=True)
-if __name__ == '__main__':
+if __name__ == "__main__":
n = 0
- for file_path in tqdm(list(ROBOTS_VIEWS_DIR.glob('*.jpg')), unit='image', desc='Moving aerial views'):
+ for file_path in tqdm(list(ROBOTS_VIEWS_DIR.glob("*.jpg")), unit="image", desc="Moving aerial views"):
if aerial_view_detector.is_matching(io.imread(str(file_path))):
move(str(file_path), str(AERIAL_VIEWS_DIR / file_path.name))
n += 1
- print(f'Moved {n} images')
+ print(f"Moved {n} images")
diff --git a/drone-at-base/research/constants.py b/drone-at-base/research/constants.py
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..8b137891791fe96927ad78e64b0aad7bded08bdc 100644
--- a/drone-at-base/research/constants.py
+++ b/drone-at-base/research/constants.py
@@ -0,0 +1 @@
+