diff --git a/.gitignore b/.gitignore
index 95615a2a008170b9320e24a7b0ec0d8f5e1e55b8..6ad15ab1ede0cbec7aa83b0963be13732409c647 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,6 +5,7 @@
# TF models
/models
+do_not_commit.py
# Byte-compiled / optimized / DLL files
__pycache__/
diff --git a/common/poetry.lock b/common/poetry.lock
index 696dcb36d509498476ab2f5a1b67451f12f95f40..9a42637dd95b439a257bc035a9c81a7198ac2003 100644
Binary files a/common/poetry.lock and b/common/poetry.lock differ
diff --git a/common/polystar/common/frame_generators/__init__.py b/common/polystar/common/frame_generators/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/common/polystar/common/frame_generators/camera_frame_generator.py b/common/polystar/common/frame_generators/camera_frame_generator.py
new file mode 100644
index 0000000000000000000000000000000000000000..e613cb6ea4796c6faf7d90a4899a19eb482e3dfe
--- /dev/null
+++ b/common/polystar/common/frame_generators/camera_frame_generator.py
@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+from typing import Any, Iterable
+
+import cv2
+
+from polystar.common.frame_generators.cv2_frame_generator_abc import CV2FrameGeneratorABC
+
+
+@dataclass
+class CameraFrameGenerator(CV2FrameGeneratorABC):
+ width: int
+ height: int
+
+ def _capture_params(self) -> Iterable[Any]:
+ return (
+ "nvarguscamerasrc ! "
+ "video/x-raw(memory:NVMM), "
+ f"width=(int){self.width}, height=(int){self.height}, "
+ "format=(string)NV12, framerate=(fraction)60/1 ! "
+ "nvvidconv flip-method=0 ! "
+ f"video/x-raw, width=(int){self.width}, height=(int){self.height}, "
+ "format=(string)BGRx ! "
+ "videoconvert ! appsink",
+ cv2.CAP_GSTREAMER,
+ )
diff --git a/common/polystar/common/frame_generators/cv2_frame_generator_abc.py b/common/polystar/common/frame_generators/cv2_frame_generator_abc.py
new file mode 100644
index 0000000000000000000000000000000000000000..7cdf709cfdfe44732494be77a70bdcdb57948aca
--- /dev/null
+++ b/common/polystar/common/frame_generators/cv2_frame_generator_abc.py
@@ -0,0 +1,32 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Iterable
+
+import cv2
+
+from polystar.common.frame_generators.frames_generator_abc import FrameGeneratorABC
+from polystar.common.models.image import Image
+
+
+@dataclass
+class CV2FrameGeneratorABC(FrameGeneratorABC, ABC):
+
+ _cap: cv2.VideoCapture = field(init=False, repr=False)
+
+ def __enter__(self):
+ self._cap = cv2.VideoCapture(*self._capture_params())
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self._cap.release()
+
+ def generate(self) -> Iterable[Image]:
+ with self:
+ while 1:
+ is_open, frame = self._cap.read()
+ if not is_open:
+ return
+ yield frame
+
+ @abstractmethod
+ def _capture_params(self) -> Iterable[Any]:
+ pass
diff --git a/common/polystar/common/frame_generators/fps_video_frame_generator.py b/common/polystar/common/frame_generators/fps_video_frame_generator.py
new file mode 100644
index 0000000000000000000000000000000000000000..ade93a8882e9259e4fb23178eea38a01bc777df5
--- /dev/null
+++ b/common/polystar/common/frame_generators/fps_video_frame_generator.py
@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import Iterable
+
+import ffmpeg
+
+from polystar.common.frame_generators.video_frame_generator import VideoFrameGenerator
+from polystar.common.models.image import Image
+
+
+@dataclass
+class FPSVideoFrameGenerator(VideoFrameGenerator):
+
+ desired_fps: int
+
+ def __post_init__(self):
+ self.frame_rate: int = self._get_video_fps() // self.desired_fps
+
+ def _get_video_fps(self):
+ return max(
+ int(stream["r_frame_rate"].split("/")[0]) for stream in ffmpeg.probe(str(self.video_path))["streams"]
+ )
+
+ def generate(self) -> Iterable[Image]:
+ for i, frame in enumerate(super().generate()):
+ if not i % self.frame_rate:
+ yield frame
diff --git a/common/polystar/common/frame_generators/frames_generator_abc.py b/common/polystar/common/frame_generators/frames_generator_abc.py
new file mode 100644
index 0000000000000000000000000000000000000000..737446b257d3eef1971a6b1d1b1e242d8bbc2726
--- /dev/null
+++ b/common/polystar/common/frame_generators/frames_generator_abc.py
@@ -0,0 +1,10 @@
+from abc import ABC, abstractmethod
+from typing import Iterable
+
+from polystar.common.models.image import Image
+
+
+class FrameGeneratorABC(ABC):
+ @abstractmethod
+ def generate(self) -> Iterable[Image]:
+ pass
diff --git a/common/polystar/common/frame_generators/video_frame_generator.py b/common/polystar/common/frame_generators/video_frame_generator.py
new file mode 100644
index 0000000000000000000000000000000000000000..388dcc2a6999d27b14908c79c33fe9a30f345eec
--- /dev/null
+++ b/common/polystar/common/frame_generators/video_frame_generator.py
@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, Any
+
+from polystar.common.frame_generators.cv2_frame_generator_abc import CV2FrameGeneratorABC
+
+
+@dataclass
+class VideoFrameGenerator(CV2FrameGeneratorABC):
+
+ video_path: Path
+
+ def _capture_params(self) -> Iterable[Any]:
+ return (str(self.video_path),)
diff --git a/common/polystar/common/models/trt_model.py b/common/polystar/common/models/trt_model.py
index 676e06af75a83aad7058d873874c848a73cc743f..b72fa2020009e6c03ba1fb964cdb62cdbdcd0489 100644
--- a/common/polystar/common/models/trt_model.py
+++ b/common/polystar/common/models/trt_model.py
@@ -46,7 +46,7 @@ class TRTModel:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, self.input_size)
img = img.transpose((2, 0, 1)).astype(np.float32)
- img = (2.0/255.0) * img - 1.0
+ img = (2.0 / 255.0) * img - 1.0
return img
# Initialization
diff --git a/common/polystar/common/utils/video_frame_generator.py b/common/polystar/common/utils/video_frame_generator.py
deleted file mode 100644
index e28e7b7416adb0e79b5d70ea62492eb93a30d34c..0000000000000000000000000000000000000000
--- a/common/polystar/common/utils/video_frame_generator.py
+++ /dev/null
@@ -1,29 +0,0 @@
-from pathlib import Path
-
-import cv2
-import ffmpeg
-
-
-class VideoFrameGenerator:
- def __init__(self, video_path: Path, desired_fps: int):
- self.video_path: Path = video_path
- self.desired_fps: int = desired_fps
- self.video_fps: int = self._get_video_fps()
-
- def _get_video_fps(self):
- return max(
- int(stream["r_frame_rate"].split("/")[0]) for stream in ffmpeg.probe(str(self.video_path))["streams"]
- )
-
- def generate(self):
- video = cv2.VideoCapture(str(self.video_path))
- frame_rate = self.video_fps // self.desired_fps
- count = 0
- while 1:
- is_unfinished, frame = video.read()
- if not is_unfinished:
- video.release()
- return
- if not count % frame_rate:
- yield frame
- count += 1
diff --git a/common/polystar/common/view/bend_object_on_image.py b/common/polystar/common/view/bend_object_on_image.py
deleted file mode 100644
index b91ce7206d7e4f2e6514bccf60ee66f23b9f2690..0000000000000000000000000000000000000000
--- a/common/polystar/common/view/bend_object_on_image.py
+++ /dev/null
@@ -1,58 +0,0 @@
-import cv2
-from typing import Tuple
-
-import numpy as np
-
-from polystar.common.models.image import Image
-from polystar.common.models.object import Object
-
-_COLORS = [
- [31, 119, 180],
- [255, 127, 14],
- [44, 160, 44],
- [214, 39, 40],
- [148, 103, 189],
- [140, 86, 75],
- [227, 119, 194],
- [127, 127, 127],
- [188, 189, 34],
- [23, 190, 207],
-] # seaborn.color_palette() * 255
-
-ALPHA = 0.5
-FONT = cv2.FONT_HERSHEY_PLAIN
-TEXT_SCALE = 1.0
-TEXT_THICKNESS = 1
-BLACK = (0, 0, 0)
-WHITE = (255, 255, 255)
-
-
-def bend_boxed_text_on_image(img: Image, text: str, topleft: Tuple[int, int], color: Tuple[int, int, int]):
- assert img.dtype == np.uint8
- img_h, img_w, _ = img.shape
- if topleft[0] >= img_w or topleft[1] >= img_h:
- return img
- margin = 3
- size = cv2.getTextSize(text, FONT, TEXT_SCALE, TEXT_THICKNESS)
- w = size[0][0] + margin * 2
- h = size[0][1] + margin * 2
- # the patch is used to draw boxed text
- patch = np.zeros((h, w, 3), dtype=np.uint8)
- patch[...] = color
- cv2.putText(patch, text, (margin+1, h-margin-2), FONT, TEXT_SCALE,
- WHITE, thickness=TEXT_THICKNESS, lineType=cv2.LINE_8)
- cv2.rectangle(patch, (0, 0), (w-1, h-1), BLACK, thickness=1)
- w = min(w, img_w - topleft[0]) # clip overlay at image boundary
- h = min(h, img_h - topleft[1])
- # Overlay the boxed text onto region of interest (roi) in img
- roi = img[topleft[1]:topleft[1]+h, topleft[0]:topleft[0]+w, :]
- cv2.addWeighted(patch[0:h, 0:w, :], ALPHA, roi, 1 - ALPHA, 0, roi)
- return img
-
-
-def bend_object_on_image(image: Image, obj: Object):
- color = _COLORS[obj.type.value]
- cv2.rectangle(image, (obj.x, obj.y), (obj.x + obj.w, obj.y + obj.h), color, 2)
-
- bend_boxed_text_on_image(image, f"{obj.type.name} ({obj.confidence:.1%})", (obj.x, obj.y), _COLORS[obj.type.value])
- return image
diff --git a/common/polystar/common/view/cv2_results_viewer.py b/common/polystar/common/view/cv2_results_viewer.py
new file mode 100644
index 0000000000000000000000000000000000000000..3894982d9c1ee40a93830b744fdd76e99826b222
--- /dev/null
+++ b/common/polystar/common/view/cv2_results_viewer.py
@@ -0,0 +1,73 @@
+import cv2
+import numpy as np
+
+from polystar.common.models.image import Image
+from polystar.common.view.results_viewer_abc import ResultViewerABC, ColorView
+
+ALPHA = 0.5
+FONT = cv2.FONT_HERSHEY_PLAIN
+TEXT_SCALE = 1.0
+TEXT_THICKNESS = 1
+BLACK = (0, 0, 0)
+WHITE = (255, 255, 255)
+COLORS = [
+ (31, 119, 180),
+ (255, 127, 14),
+ (44, 160, 44),
+ (214, 39, 40),
+ (148, 103, 189),
+ (140, 86, 75),
+ (227, 119, 194),
+ (127, 127, 127),
+ (188, 189, 34),
+ (23, 190, 207),
+] # seaborn.color_palette() * 255
+
+
+class CV2ResultViewer(ResultViewerABC):
+ def __init__(self, name: str, delay: int = 1, end_keys: str = "q"):
+ self.end_keys = [ord(c) for c in end_keys]
+ self.delay = delay
+ self.name = name
+ self._current_image: Image = None
+ self.finished = False
+ super().__init__(COLORS)
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ cv2.destroyWindow(self.name)
+
+ def new(self, image: Image):
+ self.height, self.width, _ = image.shape
+ self._current_image = image
+
+ def add_text(self, text: str, x: int, y: int, color: ColorView):
+ margin = 3
+ size = cv2.getTextSize(text, FONT, TEXT_SCALE, TEXT_THICKNESS)
+ w = size[0][0] + margin * 2
+ h = size[0][1] + margin * 2
+ # the patch is used to draw boxed text
+ patch = np.zeros((h, w, 3), dtype=np.uint8)
+ patch[...] = color
+ cv2.putText(
+ patch,
+ text,
+ (margin + 1, h - margin - 2),
+ FONT,
+ TEXT_SCALE,
+ WHITE,
+ thickness=TEXT_THICKNESS,
+ lineType=cv2.LINE_8,
+ )
+ cv2.rectangle(patch, (0, 0), (w - 1, h - 1), BLACK, thickness=1)
+ w = min(w, self.width - x) # clip overlay at image boundary
+ h = min(h, self.height - y)
+ # Overlay the boxed text onto region of interest (roi) in img
+ roi = self._current_image[y : y + h, x : x + w, :]
+ cv2.addWeighted(patch[0:h, 0:w, :], ALPHA, roi, 1 - ALPHA, 0, roi)
+
+ def add_rectangle(self, x: int, y: int, w: int, h: int, color: ColorView):
+ cv2.rectangle(self._current_image, (x, y), (x + w, y + h), color, 2)
+
+ def display(self):
+ cv2.imshow(self.name, self._current_image)
+ self.finished = cv2.waitKey(self.delay) & 0xFF in self.end_keys
diff --git a/common/polystar/common/view/plt_display_image_with_annotation.py b/common/polystar/common/view/plt_display_image_with_annotation.py
deleted file mode 100644
index 8c28081a8bf2569ba08c2ca3b999618bfe34285f..0000000000000000000000000000000000000000
--- a/common/polystar/common/view/plt_display_image_with_annotation.py
+++ /dev/null
@@ -1,29 +0,0 @@
-from typing import Iterable
-
-import matplotlib.pyplot as plt
-
-from polystar.common.models.image import Image
-from polystar.common.models.image_annotation import ImageAnnotation
-from polystar.common.models.object import Object
-import seaborn as sns
-
-
-_COLORS = sns.color_palette()
-
-
-def display_image_with_objects(image: Image, objects: Iterable[Object]):
- plt.figure(figsize=(16, 9))
- plt.imshow(image)
- for obj in objects:
- if obj.confidence >= 0.5:
- color = _COLORS[obj.type.value]
- rect = plt.Rectangle((obj.x, obj.y), obj.w, obj.h, linewidth=2, edgecolor=color, fill=False)
- plt.gca().add_patch(rect)
- plt.text(obj.x, obj.y - 2, f"{obj.type.name} ({int(obj.confidence*100)}%)", color=color, weight="bold")
- plt.axis("off")
- plt.tight_layout()
- plt.show()
-
-
-def display_image_annotation(image_annotation: ImageAnnotation):
- display_image_with_objects(image_annotation.image, image_annotation.objects)
diff --git a/common/polystar/common/view/plt_display_image_with_object.py b/common/polystar/common/view/plt_display_image_with_object.py
deleted file mode 100644
index 075bc85d9812b3c00be55e8a660e39869434cad6..0000000000000000000000000000000000000000
--- a/common/polystar/common/view/plt_display_image_with_object.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from polystar.common.models.image import Image
-from polystar.common.models.object import Object
-from polystar.common.view.plt_display_image_with_annotation import display_image_with_objects
-
-
-def display_object(image: Image, obj: Object):
- display_image_with_objects(image, [obj])
diff --git a/common/polystar/common/view/plt_results_viewer.py b/common/polystar/common/view/plt_results_viewer.py
new file mode 100644
index 0000000000000000000000000000000000000000..e46e019edb1eeccf8bc32da9901f84362c5e5aa5
--- /dev/null
+++ b/common/polystar/common/view/plt_results_viewer.py
@@ -0,0 +1,46 @@
+from typing import Tuple
+
+import matplotlib.pyplot as plt
+
+from polystar.common.models.image import Image
+from polystar.common.view.results_viewer_abc import ResultViewerABC, ColorView
+
+
+COLORS = [
+ (0.12156862745098039, 0.4666666666666667, 0.7058823529411765),
+ (1.0, 0.4980392156862745, 0.054901960784313725),
+ (0.17254901960784313, 0.6274509803921569, 0.17254901960784313),
+ (0.8392156862745098, 0.15294117647058825, 0.1568627450980392),
+ (0.5803921568627451, 0.403921568627451, 0.7411764705882353),
+ (0.5490196078431373, 0.33725490196078434, 0.29411764705882354),
+ (0.8901960784313725, 0.4666666666666667, 0.7607843137254902),
+ (0.4980392156862745, 0.4980392156862745, 0.4980392156862745),
+ (0.7372549019607844, 0.7411764705882353, 0.13333333333333333),
+ (0.09019607843137255, 0.7450980392156863, 0.8117647058823529),
+]
+
+
+class PltResultViewer(ResultViewerABC):
+ def __init__(self, name: str, fig_size: Tuple[int, int] = (16, 9)):
+ self.name = name
+ self.fig_size = fig_size
+ self._current_fig = None
+ super().__init__(COLORS)
+
+ def new(self, image: Image):
+ self._current_fig = plt.figure(figsize=self.fig_size)
+ plt.imshow(image)
+ plt.title(self.name)
+ plt.axis("off")
+
+ def add_text(self, text: str, x: int, y: int, color: ColorView):
+ plt.text(x, y - 2, text, color=color, weight="bold")
+
+ def add_rectangle(self, x: int, y: int, w: int, h: int, color: ColorView):
+ rect = plt.Rectangle((x, y), w, h, linewidth=2, edgecolor=color, fill=False)
+ plt.gca().add_patch(rect)
+
+ def display(self):
+ plt.tight_layout()
+ plt.show()
+ plt.close(self._current_fig)
diff --git a/common/polystar/common/view/results_viewer_abc.py b/common/polystar/common/view/results_viewer_abc.py
new file mode 100644
index 0000000000000000000000000000000000000000..486bf0d29bb288453b158acde5ef0d95e084f974
--- /dev/null
+++ b/common/polystar/common/view/results_viewer_abc.py
@@ -0,0 +1,52 @@
+from abc import ABC, abstractmethod
+from typing import Iterable, Tuple, NewType, Sequence
+
+from polystar.common.models.image import Image
+from polystar.common.models.image_annotation import ImageAnnotation
+from polystar.common.models.object import Object
+
+ColorView = NewType("ColorView", Tuple[float, float, float])
+
+
+class ResultViewerABC(ABC):
+ def __init__(self, colors: Sequence[ColorView]):
+ self.colors = colors
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ pass
+
+ @abstractmethod
+ def new(self, image: Image):
+ pass
+
+ @abstractmethod
+ def add_text(self, text: str, x: int, y: int, color: ColorView):
+ pass
+
+ @abstractmethod
+ def add_rectangle(self, x: int, y: int, w: int, h: int, color: ColorView):
+ pass
+
+ @abstractmethod
+ def display(self):
+ pass
+
+ def add_object(self, obj: Object):
+ color = self.colors[obj.type.value]
+ self.add_rectangle(obj.x, obj.y, obj.w, obj.h, color)
+ self.add_text(f"{obj.type.name} ({obj.confidence:.1%})", obj.x, obj.y, color)
+
+ def add_objects(self, objects: Iterable[Object]):
+ for obj in objects:
+ self.add_object(obj)
+
+ def display_image_with_objects(self, image: Image, objects: Iterable[Object]):
+ self.new(image)
+ self.add_objects(objects)
+ self.display()
+
+ def display_image_annotation(self, annotation: ImageAnnotation):
+ self.display_image_with_objects(annotation.image, annotation.objects)
diff --git a/common/pyproject.toml b/common/pyproject.toml
index d499fd72844b1516a60a602a1c12277f167151e7..9590454498d87945cedddbc85223986fa9531c8b 100644
--- a/common/pyproject.toml
+++ b/common/pyproject.toml
@@ -8,7 +8,7 @@ packages = [{ include = "polystar" }]
[tool.poetry.dependencies]
python = "^3.7"
dynaconf = "^2.2.2"
-opencv-python = "^4.2.0"
+opencv-python = "4.1.x"
injector = "^0.18.3"
numpy = "1.18.x"
matplotlib = "^3.1.3"
diff --git a/common/research_common/dataset/twitch/robots_views_extractor.py b/common/research_common/dataset/twitch/robots_views_extractor.py
index 878295d89be72f7eab47cf099b829c5d48a90294..c14759bc74ab4250992d3d5f5a6e204c1dedeacc 100644
--- a/common/research_common/dataset/twitch/robots_views_extractor.py
+++ b/common/research_common/dataset/twitch/robots_views_extractor.py
@@ -3,7 +3,7 @@ import ffmpeg
import numpy as np
from tqdm import tqdm
-from polystar.common.utils.video_frame_generator import VideoFrameGenerator
+from polystar.common.frame_generators.fps_video_frame_generator import FPSVideoFrameGenerator
from research_common.constants import TWITCH_DSET_DIR, TWITCH_ROBOTS_VIEWS_DIR
from research_common.dataset.twitch.mask_detector import is_image_from_robot_view
@@ -15,7 +15,7 @@ class RobotsViewExtractor:
def __init__(self, video_name: str):
self.video_name: str = video_name
self.video_path = TWITCH_DSET_DIR / "videos" / f"{video_name}.mp4"
- self.frame_generator: VideoFrameGenerator = VideoFrameGenerator(self.video_path, self.FPS)
+ self.frame_generator: FPSVideoFrameGenerator = FPSVideoFrameGenerator(self.video_path, self.FPS)
self.count = 0
(TWITCH_ROBOTS_VIEWS_DIR / self.video_name).mkdir(exist_ok=True)
diff --git a/dataset/twitch/runes/.gitignore b/dataset/twitch/runes/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..c96a04f008ee21e260b28f7701595ed59e2839e3
--- /dev/null
+++ b/dataset/twitch/runes/.gitignore
@@ -0,0 +1,2 @@
+*
+!.gitignore
\ No newline at end of file
diff --git a/drone-at-base/poetry.lock b/drone-at-base/poetry.lock
index a71cb44e19bfb263ef690b2f061b8284fbf17e53..5d5ab9b3fe409ee080089ae7f272c78a9f47c25b 100644
Binary files a/drone-at-base/poetry.lock and b/drone-at-base/poetry.lock differ
diff --git a/models-training/install_tensorflow_models.sh b/models-training/install_tensorflow_models.sh
deleted file mode 100755
index 80e4d45aad2647a54ca637c58b2aa13c93f0b59b..0000000000000000000000000000000000000000
--- a/models-training/install_tensorflow_models.sh
+++ /dev/null
@@ -1,52 +0,0 @@
-#!/bin/bash
-
-ROOT_DIR=`pwd`
-MODELS_DIR=$ROOT_DIR/../models
-PYTHON='poetry run python'
-PIP='poetry run pip'
-
-# make sure tensorflow has been installed
-$PIP list | grep tensorflow
-if [ $? -ne 0 ]; then
- echo "TensorFlow doesn't seem to be installed!"
- exit
-fi
-
-# Download tensorflow models, with the working commit
-cd ..
-git clone https://github.com/tensorflow/models
-cd $MODELS_DIR || exit
-git checkout 6518c1c
-
-# Fix the code for python3
-cd research || exit
-sed -i "" -e "157s/print '--annotation_type expected value is 1 or 2.'/print('--annotation_type expected value is 1 or 2.')/" \
- object_detection/dataset_tools/oid_hierarchical_labels_expansion.py
-sed -i "" -e "516s/print num_classes, num_anchors/print(num_classes, num_anchors)/" \
- object_detection/meta_architectures/ssd_meta_arch_test.py
-sed -i "" -e "282s/losses_dict.itervalues()/losses_dict.values()/" \
- object_detection/model_lib.py
-sed -i "" -e "381s/category_index.values(),/list(category_index.values()),/" \
- object_detection/model_lib.py
-sed -i "" -e "391s/eval_metric_ops.iteritems()/eval_metric_ops.items()/" \
- object_detection/model_lib.py
-sed -i "" -e "225s/reversed(zip(output_feature_map_keys, output_feature_maps_list)))/reversed(list(zip(output_feature_map_keys, output_feature_maps_list))))/" \
- object_detection/models/feature_map_generators.py
-sed -i "" -e "842s/print 'Scores and tpfp per class label: {}'.format(class_index)/print('Scores and tpfp per class label: {}'.format(class_index))/" \
- object_detection/utils/object_detection_evaluation.py
-sed -i "" -e "843s/print tp_fp_labels/print(tp_fp_labels)/" \
- object_detection/utils/object_detection_evaluation.py
-sed -i "" -e "844s/print scores/print(scores)/" \
- object_detection/utils/object_detection_evaluation.py
-sed -n '31p' object_detection/eval_util.py | grep -q vis_utils &&
- ex -s -c 31m23 -c w -c q object_detection/eval_util.py
-
-protoc object_detection/protos/*.proto --python_out=.
-
-# run a basic test to make sure tensorflow object detection is working
-echo Running model_builder_test.py
-CUDA_VISIBLE_DEVICES=0 \
-PYTHONPATH=$MODELS_DIR/research:$MODELS_DIR/research/slim \
- $PYTHON $MODELS_DIR/research/object_detection/builders/model_builder_test.py
-
-echo add tf models, research and slim to your PYTHONPATH
diff --git a/models-training/poetry.lock b/models-training/poetry.lock
deleted file mode 100644
index 76d867b2fb54391a730e78314f4dcf473344f33b..0000000000000000000000000000000000000000
Binary files a/models-training/poetry.lock and /dev/null differ
diff --git a/models-training/pyproject.toml b/models-training/pyproject.toml
deleted file mode 100644
index eea43767a4de8e3691c58cde5764b686abbcae48..0000000000000000000000000000000000000000
--- a/models-training/pyproject.toml
+++ /dev/null
@@ -1,29 +0,0 @@
-[tool.poetry]
-name = "polystar.models-training"
-version = "0.1.0"
-description = ""
-authors = ["Polystar"]
-packages = [{ include = "polystar" }]
-
-[tool.poetry.dependencies]
-python = "3.6"
-#tensorflow-gpu = { version = "1.12.x", platform = "linux/win32" }
-matplotlib = "^3.1.3"
-PILLOW = "^7.0.0"
-Cython = "^0.29.15"
-pycocotools = "^2.0.0"
-six = "^1.14.0"
-numpy = "^1.18.1"
-wheel = "^0.34.2"
-setuptools = "^45.2.0"
-mock = "^4.0.1"
-future = "^0.18.2"
-keras_applications = "^1.0.8"
-keras_preprocessing = "^1.1.0"
-
-[tool.poetry.dev-dependencies]
-pytest = "^4.5"
-
-[tool.black]
-line-length = 120
-target_version = ['py36']
diff --git a/robots-at-robots/poetry.lock b/robots-at-robots/poetry.lock
index a71cb44e19bfb263ef690b2f061b8284fbf17e53..5d5ab9b3fe409ee080089ae7f272c78a9f47c25b 100644
Binary files a/robots-at-robots/poetry.lock and b/robots-at-robots/poetry.lock differ
diff --git a/robots-at-robots/research/demos/demo_infer.py b/robots-at-robots/research/demos/demo_infer.py
index 4e0a7f237d67847d916a9318e520897fdea3fe2c..9769460c3cf395f0d1729138d17683c3695bbc46 100644
--- a/robots-at-robots/research/demos/demo_infer.py
+++ b/robots-at-robots/research/demos/demo_infer.py
@@ -2,7 +2,7 @@ from polystar.common.models.label_map import LabelMap
from polystar.common.pipeline.objects_detectors.tf_model_objects_detector import TFModelObjectsDetector
from polystar.common.pipeline.objects_validators.confidence_object_validator import ConfidenceObjectValidator
from polystar.common.utils.tensorflow import patch_tf_v2
-from polystar.common.view.plt_display_image_with_annotation import display_image_with_objects
+from polystar.common.view.plt_results_viewer import PltResultViewer
from polystar.robots_at_robots.dependency_injection import make_injector
from research.demos.utils import load_tf_model
from research_common.dataset.dji.dji_roco_datasets import DJIROCODataset
@@ -16,11 +16,13 @@ if __name__ == "__main__":
objects_detector = TFModelObjectsDetector(load_tf_model(), injector.get(LabelMap))
filters = [ConfidenceObjectValidator(confidence_threshold=0.5)]
- for i, image in enumerate(SplitDataset(DJIROCODataset.CentralChina, Split.Test).images):
- objects = objects_detector.detect(image)
- for f in filters:
- objects = f.filter(objects, image)
+ with PltResultViewer("Demo of tf model") as viewer:
+ for i, image in enumerate(SplitDataset(DJIROCODataset.CentralChina, Split.Test).images):
+ objects = objects_detector.detect(image)
+ for f in filters:
+ objects = f.filter(objects, image)
- display_image_with_objects(image, objects)
- if i == 0:
- break
+ viewer.display_image_with_objects(image, objects)
+
+ if i == 5:
+ break
diff --git a/robots-at-robots/research/demos/demo_pipeline.py b/robots-at-robots/research/demos/demo_pipeline.py
index 047aaf2a7d0c120fb5f6dd9e81ffb3636e43a918..ca58ed5fbc583e81a082ff747036caa7ea3a4a00 100644
--- a/robots-at-robots/research/demos/demo_pipeline.py
+++ b/robots-at-robots/research/demos/demo_pipeline.py
@@ -10,7 +10,7 @@ from polystar.common.pipeline.objects_validators.type_object_validator import Ty
from polystar.common.pipeline.pipeline import Pipeline
from polystar.common.pipeline.target_factories.ratio_simple_target_factory import RatioSimpleTargetFactory
from polystar.common.utils.tensorflow import patch_tf_v2
-from polystar.common.view.plt_display_image_with_object import display_object
+from polystar.common.view.plt_results_viewer import PltResultViewer
from polystar.robots_at_robots.dependency_injection import make_injector
from research.demos.utils import load_tf_model
from research_common.dataset.dji.dji_roco_datasets import DJIROCODataset
@@ -28,11 +28,12 @@ if __name__ == "__main__":
target_factory=RatioSimpleTargetFactory(injector.get(Camera), 300, 100),
)
- for i, image_path in enumerate(SplitDataset(DJIROCODataset.CentralChina, Split.Test).image_paths):
- image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
- obj = pipeline.predict_best_object(image)
+ with PltResultViewer("Demo of tf model") as viewer:
+ for i, image_path in enumerate(SplitDataset(DJIROCODataset.CentralChina, Split.Test).image_paths):
+ image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
+ obj = pipeline.predict_best_object(image)
- display_object(image, obj)
+ viewer.display_image_with_objects(image, [obj])
- if i == 0:
- break
+ if i == 5:
+ break
diff --git a/robots-at-robots/research/demos/demo_pipeline_camera.py b/robots-at-robots/research/demos/demo_pipeline_camera.py
index 94b8bc8417664e606080264e1ab2928241c9df76..33f15701452d229d2e76b445bf0ed79d507964db 100644
--- a/robots-at-robots/research/demos/demo_pipeline_camera.py
+++ b/robots-at-robots/research/demos/demo_pipeline_camera.py
@@ -1,59 +1,21 @@
-import subprocess
-import sys
-
-import cv2
from time import time
import pycuda.autoinit # This is needed for initializing CUDA driver
from polystar.common.constants import MODELS_DIR
+from polystar.common.frame_generators.camera_frame_generator import CameraFrameGenerator
from polystar.common.models.label_map import LabelMap
from polystar.common.models.trt_model import TRTModel
from polystar.common.pipeline.objects_detectors.trt_model_object_detector import TRTModelObjectsDetector
from polystar.common.pipeline.objects_validators.confidence_object_validator import ConfidenceObjectValidator
from polystar.common.utils.tensorflow import patch_tf_v2
-from polystar.common.view.bend_object_on_image import bend_object_on_image, bend_boxed_text_on_image
+from polystar.common.view.cv2_results_viewer import CV2ResultViewer
from polystar.robots_at_robots.dependency_injection import make_injector
from polystar.robots_at_robots.globals import settings
-WINDOWS_NAME = "TensorRT demo"
-
-
[pycuda.autoinit] # So pycharm won't remove the import
-def open_cam_onboard(width, height):
- """Open the Jetson onboard camera."""
- gst_elements = str(subprocess.check_output("gst-inspect-1.0"))
- if "nvcamerasrc" in gst_elements:
- # On versions of L4T prior to 28.1, you might need to add
- # 'flip-method=2' into gst_str below.
- gst_str = (
- "nvcamerasrc ! "
- "video/x-raw(memory:NVMM), "
- "width=(int)2592, height=(int)1458, "
- "format=(string)I420, framerate=(fraction)30/1 ! "
- "nvvidconv ! "
- "video/x-raw, width=(int){}, height=(int){}, "
- "format=(string)BGRx ! "
- "videoconvert ! appsink"
- ).format(width, height)
- elif "nvarguscamerasrc" in gst_elements:
- gst_str = (
- "nvarguscamerasrc ! "
- "video/x-raw(memory:NVMM), "
- f"width=(int){width}, height=(int){height}, "
- "format=(string)NV12, framerate=(fraction)60/1 ! "
- "nvvidconv flip-method=0 ! "
- f"video/x-raw, width=(int){width}, height=(int){height}, "
- "format=(string)BGRx ! "
- "videoconvert ! appsink"
- )
- else:
- raise RuntimeError("onboard camera source not found!")
- return cv2.VideoCapture(gst_str, cv2.CAP_GSTREAMER)
-
-
if __name__ == "__main__":
patch_tf_v2()
injector = make_injector()
@@ -63,31 +25,22 @@ if __name__ == "__main__":
)
filters = [ConfidenceObjectValidator(confidence_threshold=0.5)]
- cap = open_cam_onboard(1_280, 720)
-
- if not cap.isOpened():
- sys.exit("Failed to open camera!")
-
fps = 0
- try:
- while True:
+ with CV2ResultViewer("TensorRT demo") as viewer:
+ for image in CameraFrameGenerator(1_280, 720).generate():
+
previous_time = time()
- ret, image = cap.read()
+
+ # inference
objects = objects_detector.detect(image)
for f in filters:
objects = f.filter(objects, image)
- fps = .9 * fps + .1 / (time() - previous_time)
- bend_boxed_text_on_image(image, f'FPS: {fps:.1f}', (10, 10), (0, 0, 0))
-
- for obj in objects:
- bend_object_on_image(image, obj)
-
- # Display the resulting frame
- cv2.imshow("frame", image)
- if cv2.waitKey(1) & 0xFF == ord("q"):
+ # display
+ fps = 0.9 * fps + 0.1 / (time() - previous_time)
+ viewer.new(image)
+ viewer.add_objects(objects)
+ viewer.add_text(f"FPS: {fps:.1f}", 10, 10, (0, 0, 0))
+ viewer.display()
+ if viewer.finished:
break
- finally:
- # When everything done, release the capture
- cap.release()
- cv2.destroyAllWindows()
diff --git a/robots-at-robots/research/demos/utils.py b/robots-at-robots/research/demos/utils.py
index 93969a53989b036dd763f68b07222bb4522c37a0..6bcbcdba5703dfcfe3dea7e808258dc760555d7f 100644
--- a/robots-at-robots/research/demos/utils.py
+++ b/robots-at-robots/research/demos/utils.py
@@ -1,9 +1,8 @@
import tensorflow as tf
from polystar.common.constants import MODELS_DIR
-from polystar.robots_at_robots.globals import settings
def load_tf_model():
- model = tf.saved_model.load(export_dir=str(MODELS_DIR / settings.MODEL_NAME / "saved_model"))
+ model = tf.saved_model.load(export_dir=str(MODELS_DIR / "robots/ssd_mobilenet_v2_coco_2018_03_29" / "saved_model"))
return model.signatures["serving_default"]
diff --git a/robots-at-runes/poetry.lock b/robots-at-runes/poetry.lock
index a71cb44e19bfb263ef690b2f061b8284fbf17e53..5d5ab9b3fe409ee080089ae7f272c78a9f47c25b 100644
Binary files a/robots-at-runes/poetry.lock and b/robots-at-runes/poetry.lock differ