File size: 3,077 Bytes
beced92 a021e24 beced92 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """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)
|