def main(model_name: str = "maskrcnn_pennfudanped", score_threshold=0.55):
base_path = PROJECT_APP_PATH.user_data / 'maskrcnn'
dataset_root = Path.home() / "Data"
torch_seed(3825)
dataset = PennFudanDataset(dataset_root / "PennFudanPed",
Split.Training)
categories = dataset.categories
if True:
model = load_model(model_name=model_name,
model_directory=base_path / 'models')
else:
model = get_pretrained_instance_segmentation_maskrcnn(
dataset.response_channels)
model.to(global_torch_device())
cpu_device = torch.device("cpu")
with torch.no_grad():
with TorchEvalSession(model):
for image in tqdm(
to_tensor_generator(
frame_generator(cv2.VideoCapture(0)),
device=global_torch_device(),
)):
prediction = model(
# torch_vision_normalize_batch_nchw(
uint_hwc_to_chw_float_tensor(image).unsqueeze(0)
# )
)[0]
(boxes, labels, scores) = (
prediction["boxes"].to(cpu_device).numpy(),
prediction["labels"].to(cpu_device).numpy(),
torch.sigmoid(
prediction["scores"]).to(cpu_device).numpy(),
)
indices = scores > score_threshold
cv2.namedWindow(model_name, cv2.WINDOW_NORMAL)
cv2.imshow(
model_name,
draw_bounding_boxes(
quick_to_pil_image(image),
boxes[indices],
labels=labels[indices],
scores=scores[indices],
categories=categories,
))
if cv2.waitKey(1) == 27:
break # esc to quit
def run_webcam_demo(
cfg: NOD,
input_cfg: NOD,
categories: List,
model_ckpt: Path,
score_threshold: float = 0.7,
window_name: str = "SSD",
):
"""
:param categories:
:type categories:
:param cfg:
:type cfg:
:param model_ckpt:
:type model_ckpt:
:param score_threshold:
:type score_threshold:
:param window_name:
:type window_name:
:return:
:rtype:
"""
cpu_device = torch.device("cpu")
transforms = SSDTransform(input_cfg.image_size,
input_cfg.pixel_mean,
split=Split.Testing)
model = SingleShotDectectionNms(cfg)
checkpointer = CheckPointer(model,
save_dir=ensure_existence(
PROJECT_APP_PATH.user_data / "results"))
checkpointer.load(model_ckpt, use_latest=model_ckpt is None)
print(
f"Loaded weights from {model_ckpt if model_ckpt else checkpointer.get_checkpoint_file()}"
)
model.post_init()
model.to(global_torch_device())
with TorchEvalSession(model):
for image in tqdm(frame_generator(cv2.VideoCapture(0))):
result = model(
transforms(image)[0].unsqueeze(0).to(global_torch_device()))
height, width, *_ = image.shape
result.boxes[:, 0::2] *= width / result.img_width.cpu().item()
result.boxes[:, 1::2] *= height / result.img_height.cpu().item()
(boxes, labels, scores) = (
result.boxes.to(cpu_device).numpy(),
result.labels.to(cpu_device).numpy(),
result.scores.to(cpu_device).numpy(),
)
indices = scores > score_threshold
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.imshow(
window_name,
draw_bounding_boxes(
image,
boxes[indices],
labels=labels[indices],
scores=scores[indices],
categories=categories,
score_font=ImageFont.truetype(
PACKAGE_DATA_PATH / "Lato-Regular.ttf",
24,
),
).astype(numpy.uint8),
)
if cv2.waitKey(1) == 27:
break # esc to quit
def main():
dataset_root = Path.home() / "Data"
base_path = ensure_existence(PROJECT_APP_PATH.user_data / 'maskrcnn')
log_path = ensure_existence(PROJECT_APP_PATH.user_log / 'maskrcnn')
export_root = ensure_existence(base_path / 'models')
model_name = f'maskrcnn_pennfudanped'
batch_size = 4
num_epochs = 10
optimiser_spec = GDKC(torch.optim.Adam, lr=3e-4)
scheduler_spec = GDKC(
torch.optim.lr_scheduler.
StepLR, # a learning rate scheduler which decreases the learning rate by
step_size=3, # 10x every 3 epochs
gamma=0.1,
)
num_workers = os.cpu_count()
torch_seed(3825)
dataset = PennFudanDataset(dataset_root / "PennFudanPed",
Split.Training,
return_variant=ReturnVariant.all)
dataset_validation = PennFudanDataset(
dataset_root / "PennFudanPed",
Split.Validation,
return_variant=ReturnVariant.all,
)
split = SplitIndexer(len(dataset), validation=0.3, testing=0)
split_indices = torch.randperm(split.total_num).tolist()
data_loader = DataLoader(
Subset(dataset, split_indices[:-split.validation_num]),
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=collate_batch_fn,
)
data_loader_val = DataLoader(
Subset(dataset_validation, split_indices[-split.validation_num:]),
batch_size=1,
shuffle=False,
num_workers=num_workers,
collate_fn=collate_batch_fn,
)
model = get_pretrained_instance_segmentation_maskrcnn(
dataset.response_channels)
optimiser = optimiser_spec(trainable_parameters(model))
lr_scheduler = scheduler_spec(optimiser)
if True:
model = load_model(model_name=model_name,
model_directory=export_root)
if True:
with TorchTrainSession(model):
with TensorBoardPytorchWriter(log_path / model_name) as writer:
for epoch_i in tqdm(range(num_epochs), desc="Epoch #"):
maskrcnn_train_single_epoch(model=model,
optimiser=optimiser,
data_loader=data_loader,
writer=writer)
lr_scheduler.step() # update the learning rate
maskrcnn_evaluate(
model, data_loader_val, writer=writer
) # evaluate on the validation dataset
save_model(model,
model_name=model_name,
save_directory=export_root)
if True:
with TorchEvalSession(model): # put the model in evaluation mode
img, _ = dataset_validation[
0] # pick one image from the test set
with torch.no_grad():
prediction = model([img.to(global_torch_device())])
from matplotlib import pyplot
pyplot.imshow(
Image.fromarray(
img.mul(255).permute(1, 2, 0).byte().numpy()))
pyplot.show()
import cv2
pyplot.imshow(
Image.fromarray(prediction[0]["masks"][0, 0].mul(
255).byte().cpu().numpy()))
pyplot.show()
(boxes, labels, scores) = (
prediction[0]["boxes"].to('cpu').numpy(),
prediction[0]["labels"].to('cpu').numpy(),
torch.sigmoid(prediction[0]["scores"]).to('cpu').numpy(),
)
from draugr.opencv_utilities import draw_bounding_boxes
from draugr.torch_utilities.images.conversion import quick_to_pil_image
indices = scores > 0.1
cv2.namedWindow(model_name, cv2.WINDOW_NORMAL)
cv2.imshow(
model_name,
draw_bounding_boxes(
quick_to_pil_image(img),
boxes[indices],
labels=labels[indices],
scores=scores[indices],
#categories=categories,
))
cv2.waitKey()
def run_webcam_demo(
cfg: NOD,
categories: Sequence[str],
model_checkpoint: Path,
score_threshold: float = 0.5,
window_name: str = "SSD",
):
"""
:param categories:
:type categories:
:param cfg:
:type cfg:
:param model_checkpoint:
:type model_checkpoint:
:param score_threshold:
:type score_threshold:
:param window_name:
:type window_name:
:return:
:rtype:"""
cpu_device = torch.device("cpu")
transforms = SSDTransform(cfg.input.image_size,
cfg.input.pixel_mean,
split=SplitEnum.testing)
model = SingleShotDetection(cfg)
checkpointer = CheckPointer(model,
save_dir=ensure_existence(
PROJECT_APP_PATH.user_data / "results"))
checkpointer.load(model_checkpoint, use_latest=model_checkpoint is None)
print(
f"Loaded weights from {model_checkpoint if model_checkpoint else checkpointer.get_checkpoint_file()}"
)
model.post_init()
model.to(global_torch_device())
with TorchEvalSession(model):
for infos in tqdm(DictUnityEnvironment(connect_to_running=True)):
info = next(iter(infos.values()))
new_images = extract_all_cameras(info)
image = next(iter(new_images.values()))[..., :3][..., ::-1]
image = gamma_correct_float_to_byte(image)
result = model(
transforms(image)[0].unsqueeze(0).to(global_torch_device()))[0]
height, width, *_ = image.shape
result["boxes"][:, 0::2] *= width / result["img_width"]
result["boxes"][:, 1::2] *= height / result["img_height"]
(boxes, labels, scores) = (
result["boxes"].to(cpu_device).numpy(),
result["labels"].to(cpu_device).numpy(),
result["scores"].to(cpu_device).numpy(),
)
indices = scores > score_threshold
if show_image(
draw_bounding_boxes(
image,
boxes[indices],
labels=labels[indices],
scores=scores[indices],
categories=categories,
).astype(numpy.uint8),
window_name,
wait=1,
):
break # esc to quit
def run_traced_webcam_demo(
input_cfg: NOD,
categories: List,
score_threshold: float = 0.7,
window_name: str = "SSD",
onnx_exported: bool = False,
):
"""
:param onnx_exported:
:type onnx_exported:
:param input_cfg:
:type input_cfg:
:param categories:
:type categories:
:param score_threshold:
:type score_threshold:
:param window_name:
:type window_name:
:return:
:rtype:
"""
pass
import torch
cpu_device = torch.device("cpu")
transforms = SSDTransform(input_cfg.image_size,
input_cfg.pixel_mean,
split=Split.Testing)
model = None
if onnx_exported:
import onnx
onnx_model = onnx.load("torch_model.onnx")
onnx.checker.check_model(onnx_model)
import onnxruntime
ort_session = onnxruntime.InferenceSession("torch_model.onnx")
def to_numpy(tensor):
return (tensor.detach().cpu().numpy()
if tensor.requires_grad else tensor.cpu().numpy())
x = None
# compute onnxruntime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
else:
import torch
import io
torch.jit.load("torch_model.traced")
with open("torch_model.traced",
"rb") as f: # Load ScriptModule from io.BytesIO object
buffer = io.BytesIO(f.read())
model = torch.jit.load(
buffer) # Load all tensors to the original device
"""
buffer.seek(0)
torch.jit.load(buffer, map_location=torch.device('cpu')) # Load all tensors onto CPU, using a device
buffer.seek(0)
model = torch.jit.load(buffer, map_location='cpu') # Load all tensors onto CPU, using a string
# Load with extra files.
extra_files = torch._C.ExtraFilesMap()
extra_files['foo.txt'] = 'bar'
torch.jit.load('torch_model.traced', _extra_files=extra_files)
print(extra_files['foo.txt'])
#exit(0)
"""
with TorchDeviceSession(
device=global_torch_device(cuda_if_available=False), model=model):
with TorchEvalSession(model):
for image in tqdm(frame_generator(cv2.VideoCapture(0))):
result = SSDOut(*model(
transforms(image)[0].unsqueeze(0).to(
global_torch_device())))
height, width, *_ = image.shape
result.boxes[:, 0::2] *= width / result.img_width.cpu().item()
result.boxes[:,
1::2] *= height / result.img_height.cpu().item()
(boxes, labels, scores) = (
result.boxes.to(cpu_device).numpy(),
result.labels.to(cpu_device).numpy(),
result.scores.to(cpu_device).numpy(),
)
indices = scores > score_threshold
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
cv2.imshow(
window_name,
draw_bounding_boxes(
image,
boxes[indices],
labels=labels[indices],
scores=scores[indices],
categories=categories,
score_font=ImageFont.truetype(
PACKAGE_DATA_PATH / "Lato-Regular.ttf",
24,
),
).astype(numpy.uint8),
)
if cv2.waitKey(1) == 27:
break # esc to quit
def run_demo(cfg, class_names, model_ckpt, score_threshold, images_dir,
output_dir):
model = SingleShotDectection(cfg)
checkpointer = CheckPointer(model,
save_dir=ensure_existence(
PROJECT_APP_PATH.user_data / "results"))
checkpointer.load(model_ckpt, use_latest=model_ckpt is None)
print(
f"Loaded weights from {model_ckpt if model_ckpt else checkpointer.get_checkpoint_file()}"
)
model.post_init()
model.to(global_torch_device())
image_paths = list(images_dir.iterdir())
cpu_device = torch.device("cpu")
transforms = SSDTransform(cfg.input.image_size,
cfg.input.pixel_mean,
split=Split.Testing)
model.eval()
for i, image_path in enumerate(image_paths):
start = time.time()
image_name = os.path.basename(image_path)
image = numpy.array(Image.open(image_path).convert("RGB"))
height, width = image.shape[:2]
images = transforms(image)[0].unsqueeze(0)
load_time = time.time() - start
start = time.time()
result = model(images.to(global_torch_device()))[0]
inference_time = time.time() - start
result.boxes[:, 0::2] *= width / result.img_width
result.boxes[:, 1::2] *= height / result.img_height
(boxes, labels, scores) = (
result.boxes.to(cpu_device).numpy(),
result.labels.to(cpu_device).numpy(),
result.scores.to(cpu_device).numpy(),
)
indices = scores > score_threshold
boxes, labels, scores = boxes[indices], labels[indices], scores[
indices]
meters = " | ".join([
f"objects {len(boxes):02d}",
f"load {round(load_time * 1000):03d}ms",
f"inference {round(inference_time * 1000):03d}ms",
f"FPS {round(1.0 / inference_time)}",
])
print(f"({i + 1:04d}/{len(image_paths):04d}) {image_name}: {meters}")
drawn_image = draw_bounding_boxes(
image,
boxes,
labels,
scores,
class_names,
score_font=ImageFont.truetype(
PACKAGE_DATA_PATH / "Lato-Regular.ttf",
24,
),
).astype(numpy.uint8)
Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name))