|
|
"""Run the main code for Deep-Package-Detection""" |
|
|
|
|
|
from pathlib import Path |
|
|
import logging |
|
|
|
|
|
import click |
|
|
|
|
|
from deep_package_detection import __version__ |
|
|
from deep_package_detection.loggings import config_logger |
|
|
from deep_package_detection.detector import PackageDetectorTrainer |
|
|
from deep_package_detection.detector import PackageDetectorInference |
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
@click.group() |
|
|
@click.version_option(version=__version__) |
|
|
@click.option("-l", "--log_path", type=str, help="Path to save log file") |
|
|
@click.option( |
|
|
"-v", |
|
|
"--verbose", |
|
|
count=True, |
|
|
help="Shorthand for info/debug/warning/error loglevel (-v/-vv/-vvv/-vvvv)", |
|
|
) |
|
|
def deep_package_detection_cli(log_path: str, verbose: int) -> None: |
|
|
"""This is a deep model for detecting different types of packages.""" |
|
|
if verbose == 1: |
|
|
log_level = 10 |
|
|
elif verbose == 2: |
|
|
log_level = 20 |
|
|
elif verbose == 3: |
|
|
log_level = 30 |
|
|
else: |
|
|
log_level = 40 |
|
|
config_logger(log_level, log_path) |
|
|
|
|
|
|
|
|
@deep_package_detection_cli.command() |
|
|
@click.option("--img_resize", type=int, default=640, help="Resize images to this size.") |
|
|
@click.option( |
|
|
"--conf_path", |
|
|
type=str, |
|
|
default="src/deep_package_detection/data/data.yaml", |
|
|
help="Path to the config file", |
|
|
) |
|
|
@click.option( |
|
|
"--epochs", type=int, default=100, help="Number of epochs used in training." |
|
|
) |
|
|
@click.option("--batch_size", type=int, default=16, help="Batch size used in training.") |
|
|
@click.option( |
|
|
"--device", type=str, default="cuda", help="Use cuda or cpu for training." |
|
|
) |
|
|
def train( |
|
|
img_resize: int, conf_path: str, epochs: int, batch_size: int, device: str |
|
|
) -> None: |
|
|
"""This the CLI for training purposes""" |
|
|
detector = PackageDetectorTrainer( |
|
|
conf=conf_path, |
|
|
img_size=img_resize, |
|
|
epochs=epochs, |
|
|
device=device, |
|
|
batch_size=batch_size, |
|
|
) |
|
|
detector.train() |
|
|
_ = detector.validation() |
|
|
|
|
|
|
|
|
@deep_package_detection_cli.command() |
|
|
@click.option( |
|
|
"--model_path", |
|
|
type=click.Path(), |
|
|
default=Path("./src/deep_package_detection/model/package_detection.pt"), |
|
|
help="Path to the pre-trained model.", |
|
|
) |
|
|
@click.option( |
|
|
"--data_path", |
|
|
type=click.Path(), |
|
|
default=Path("./tests/test_data"), |
|
|
help="Path to the test data.", |
|
|
) |
|
|
@click.option( |
|
|
"--result_path", type=str, default="./results", help="Path to the results." |
|
|
) |
|
|
def infer(model_path: Path, data_path: str, result_path: str) -> None: |
|
|
"""This the CLI for testing purposes""" |
|
|
logger.info("Testing the model for package detection...") |
|
|
inferer = PackageDetectorInference( |
|
|
model_path=Path(model_path), result_path=result_path |
|
|
) |
|
|
segmentations = inferer.inference(data_path=data_path) |
|
|
counts = inferer.count_packages(segmentations) |
|
|
if counts: |
|
|
for key, val in counts.items(): |
|
|
logger.info( |
|
|
"%s packages are detected in %s as: %s", |
|
|
sum(item["count"] for item in val), |
|
|
key, |
|
|
val, |
|
|
) |
|
|
inferer.plot_and_save_results(segmentations) |
|
|
|
