add 2 m

conversion_scripts/dpt_pt_2_trt.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""
+Depth-Anything V2  (metric)
+PyTorch → ONNX → TensorRT
+
+Default: static 518×518 engine.
+Add --dynamic to build a multi-resolution engine.
+
+Error “no optimisation profile” is gone because the script
+creates a profile whenever the ONNX contains -1 dimensions.
+"""
+
+import argparse, os, sys, onnx, torch, tensorrt as trt, pycuda.driver as cuda
+cuda.init(); import pycuda.autoinit  # noqa: E402
+from depth_anything_v2.dpt import DepthAnythingV2
+
+# ───────────────────────── checkpoints ─────────────────────────
+def ckpt_path(ds, enc):
+    return f"checkpoints/depth_anything_v2_metric_{ds}_{enc}.pth"
+
+# ───────────────────────── ONNX export ─────────────────────────
+def export_onnx(a):
+    cfg = {
+        "vits": dict(encoder="vits", features=64,  out_channels=[48, 96,192,384]),
+        "vitb": dict(encoder="vitb", features=128, out_channels=[96,192,384,768]),
+        "vitl": dict(encoder="vitl", features=256, out_channels=[256,512,1024,1024])
+    }
+    print(f"[ONNX] export {a.encoder.upper()} ({a.dataset}) …")
+    model = DepthAnythingV2(**cfg[a.encoder], max_depth=a.max_depth).to("cuda")
+    model.load_state_dict(torch.load(ckpt_path(a.dataset, a.encoder), map_location="cuda"),
+                          strict=False)
+    model.eval()
+
+    dummy = torch.randn(1, 3, *a.input_hw, device="cuda")
+    dynamic_axes = {}
+    if a.dynamic:               # only add dynamic sizes if user asked for it
+        dynamic_axes = {"rgb": {0: "N", 2: "H", 3: "W"}}
+
+    torch.onnx.export(
+        model, dummy, a.onnx,
+        opset_version=17,
+        input_names=["rgb"], output_names=["depth"],
+        do_constant_folding=True, dynamic_axes=dynamic_axes)
+    onnx.checker.check_model(onnx.load(a.onnx))
+    print(f"[ONNX] saved → {a.onnx}")
+
+# ───────────────────── TensorRT build ──────────────────────────
+def build_trt(a):
+    print("[TRT] build …")
+    logger, builder = trt.Logger(trt.Logger.INFO), trt.Builder(trt.Logger(trt.Logger.INFO))
+    network = builder.create_network()
+    parser  = trt.OnnxParser(network, logger)
+    with open(a.onnx, "rb") as f:
+        if not parser.parse(f.read()):
+            for i in range(parser.num_errors):
+                print(parser.get_error(i)); sys.exit(1)
+
+    cfg = builder.create_builder_config()
+    cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, int(a.workspace_gb * (1 << 30)))
+    if a.fp16 and builder.platform_has_fast_fp16: cfg.set_flag(trt.BuilderFlag.FP16)
+    if a.sparse and hasattr(trt.BuilderFlag, "SPARSE_WEIGHTS"):
+        cfg.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
+
+    # ── add optimisation profile if input is dynamic ──
+    if network.get_input(0).shape[0] == -1:
+        prof = builder.create_optimization_profile()
+        in_name = network.get_input(0).name
+        h, w = a.input_hw
+        prof.set_shape(in_name,  (1,3,h,w), (1,3,h,w), (4,3,h,w))
+        cfg.add_optimization_profile(prof)      # mandatory :contentReference[oaicite:2]{index=2}
+
+    # TRT-10 path
+    if hasattr(builder, "build_serialized_network"):
+        eng_bytes = builder.build_serialized_network(network, cfg)
+        if eng_bytes is None: sys.exit("[ERR] build failed")
+        engine = trt.Runtime(logger).deserialize_cuda_engine(eng_bytes)
+    else:                                       # TRT-8 fallback
+        engine = builder.build_engine(network, cfg)
+
+    with open(a.engine, "wb") as f: f.write(engine.serialize())
+    print(f"[TRT] engine → {a.engine}")
+
+# ───────────────────────────── CLI ─────────────────────────────
+def parse():
+    p = argparse.ArgumentParser()
+    p.add_argument("-e","--encoder",  choices=["vits","vitb","vitl"], default="vitb")
+    p.add_argument("-d","--dataset",  choices=["vkitti","hypersim"],  default="vkitti")
+    p.add_argument("--max-depth", type=float, default=80.0)
+    p.add_argument("--input-hw",  type=int, nargs=2, default=[518,518])
+    p.add_argument("--workspace-gb", type=float, default=4.0)
+    p.add_argument("--fp16-off", dest="fp16", action="store_false")
+    p.add_argument("--sparse",   action="store_true")
+    p.add_argument("--dynamic",  action="store_true",
+                   help="make H/W & batch dynamic (adds optimisation profile)")
+    p.add_argument("--rebuild",  action="store_true")
+    return p.parse_args()
+
+def main():
+    a = parse(); a.fp16 = getattr(a,"fp16",True); a.input_hw = tuple(a.input_hw)
+    base = f"depth_anything_v2_{a.encoder}_{a.dataset}"
+    a.onnx   = f"{base}.onnx"
+    a.engine = f"{base}_{'fp16' if a.fp16 else 'fp32'}{'_dyn' if a.dynamic else ''}.engine"
+    if a.rebuild:
+        for f in (a.onnx, a.engine):
+            if os.path.isfile(f): os.remove(f)
+    if not os.path.isfile(a.onnx):  export_onnx(a)
+    if not os.path.isfile(a.engine): build_trt(a)
+    else: print(f"[OK] engine exists → {a.engine}")
+
+if __name__ == "__main__": main()
+

conversion_scripts/ultralytics-pt-to-trt.py

@@ -0,0 +1,53 @@
+import os
+import torch
+from ultralytics import YOLO
+
+# Set environment variables for CUDA
+os.environ['CUDA_HOME'] = '/usr/local/cuda'
+os.environ['PATH'] = f"{os.environ['CUDA_HOME']}/bin:{os.environ['PATH']}"
+os.environ['LD_LIBRARY_PATH'] = f"{os.environ['CUDA_HOME']}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"
+
+def verify_cuda_installation():
+    print("Verifying CUDA installation...")
+    print("CUDA_HOME:", os.environ['CUDA_HOME'])
+    print("PATH:", os.environ['PATH'])
+    print("LD_LIBRARY_PATH:", os.environ['LD_LIBRARY_PATH'])
+
+def check_cuda_availability():
+    print("Checking CUDA availability in PyTorch...")
+    cuda_available = torch.cuda.is_available()
+    print("CUDA available:", cuda_available)
+    if cuda_available:
+        device_count = torch.cuda.device_count()
+        print("CUDA device count:", device_count)
+        for i in range(device_count):
+            print(f"CUDA device {i}: {torch.cuda.get_device_name(i)}")
+    return cuda_available
+
+def load_and_transform_model(cuda_available):
+    print("Loading YOLOv8 model...")
+    model = YOLO("best_model.pt")
+    
+    if cuda_available:
+        print("CUDA is available. Exporting model to TensorRT format...")
+        model.export(format="engine")  # creates 'debris-det.engine'
+        tensorrt_model = YOLO("best.engine")
+    else:
+        print("CUDA is not available. Skipping TensorRT export.")
+        tensorrt_model = model
+    
+    return tensorrt_model
+
+def run_inference(tensorrt_model):
+    print("Running inference...")
+    results = tensorrt_model("https://ultralytics.com/images/bus.jpg")
+    
+    print("Inference results:")
+    print(results)
+
+if __name__ == "__main__":
+    verify_cuda_installation()
+    cuda_available = check_cuda_availability()
+    tensorrt_model = load_and_transform_model(cuda_available)
+    run_inference(tensorrt_model)
+

conversion_scripts/unet_to_trt_lptop.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+import os
+import torch
+import torch.nn as nn
+import tensorrt as trt
+import pycuda.autoinit  # initializes CUDA driver
+
+class Encoder(nn.Module):
+    def __init__(self, in_channels, out_channels, rate, pooling=True):
+        super(Encoder, self).__init__()
+        self.pooling = pooling
+        self.bn = nn.BatchNorm2d(in_channels)
+        self.c1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.drop = nn.Dropout(rate)
+        self.c2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(kernel_size=2)
+
+    def forward(self, x):
+        x = self.bn(x)
+        x = nn.ReLU()(self.c1(x))
+        x = self.drop(x)
+        x = nn.ReLU()(self.c2(x))
+        if self.pooling:
+            y = self.pool(x)
+            return y, x
+        return x
+
+class Decoder(nn.Module):
+    def __init__(self, in_channels, out_channels, skip_channels, rate):
+        super(Decoder, self).__init__()
+        self.bn = nn.BatchNorm2d(in_channels)
+        self.cT = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.c1 = nn.Conv2d(out_channels + skip_channels, out_channels, kernel_size=3, padding=1)
+        self.drop = nn.Dropout(rate)
+        self.c2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, x, skip_x):
+        x = self.bn(x)
+        x = nn.ReLU()(self.cT(x))
+        x = torch.cat([x, skip_x], dim=1)
+        x = nn.ReLU()(self.c1(x))
+        x = self.drop(x)
+        x = nn.ReLU()(self.c2(x))
+        return x
+
+class UNet(nn.Module):
+    def __init__(self):
+        super(UNet, self).__init__()
+        self.initial = nn.Conv2d(3, 64, kernel_size=3, padding=1)
+        self.enc1 = Encoder(64, 64, 0.1)
+        self.enc2 = Encoder(64, 128, 0.1)
+        self.enc3 = Encoder(128, 256, 0.2)
+        self.enc4 = Encoder(256, 512, 0.2)
+        self.enc5 = Encoder(512, 512, 0.3, pooling=False)
+
+        self.dec1 = Decoder(512, 512, 512, 0.2)
+        self.dec2 = Decoder(512, 256, 256, 0.2)
+        self.dec3 = Decoder(256, 128, 128, 0.1)
+        self.dec4 = Decoder(128, 64, 64, 0.1)
+
+        self.final = nn.Conv2d(64, 1, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        x = nn.ReLU()(self.initial(x))
+        p1, c1 = self.enc1(x)
+        p2, c2 = self.enc2(p1)
+        p3, c3 = self.enc3(p2)
+        p4, c4 = self.enc4(p3)
+        e = self.enc5(p4)
+
+        d1 = self.dec1(e, c4)
+        d2 = self.dec2(d1, c3)
+        d3 = self.dec3(d2, c2)
+        d4 = self.dec4(d3, c1)
+
+        return torch.sigmoid(self.final(d4))
+
+
+def convert_to_onnx(model, input_shape, onnx_path):
+    model.eval()
+    dummy_input = torch.randn(input_shape).cuda()
+    torch.onnx.export(
+        model,
+        dummy_input,
+        onnx_path,
+        verbose=False,
+        opset_version=12,
+        input_names=["input"],
+        output_names=["output"]
+    )
+    print(f"✅ ONNX model saved at: {onnx_path}")
+
+
+def convert_to_tensorrt_py(onnx_path, engine_path):
+    logger = trt.Logger(trt.Logger.WARNING)
+    builder = trt.Builder(logger)
+    network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+    network = builder.create_network(network_flags)
+    parser = trt.OnnxParser(network, logger)
+
+    with open(onnx_path, "rb") as f:
+        if not parser.parse(f.read()):
+            for i in range(parser.num_errors):
+                print(parser.get_error(i))
+            raise RuntimeError("Failed to parse ONNX model")
+
+    config = builder.create_builder_config()
+    # Set workspace memory limit to 1 GiB
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
+    if builder.platform_has_fast_fp16:
+        config.set_flag(trt.BuilderFlag.FP16)
+
+    print("⏳ Building TensorRT engine in Python...")
+    serialized_engine = builder.build_serialized_network(network, config)
+    if serialized_engine is None:
+        raise RuntimeError("Failed to build serialized engine")
+
+    # Optional: Deserialize to verify
+    runtime = trt.Runtime(logger)
+    engine = runtime.deserialize_cuda_engine(serialized_engine)
+    if engine is None:
+        raise RuntimeError("Failed to deserialize engine")
+
+    with open(engine_path, "wb") as f:
+        f.write(serialized_engine)
+        print(f"✅ TensorRT engine saved at: {engine_path}")
+
+
+def main():
+    # Paths
+    model_path = 'water_segmentation_model.pth'
+    onnx_path = 'water_segmentation_model.onnx'
+    engine_path = 'water_segmentation_model.engine'
+
+    # Load PyTorch model
+    model = UNet()
+    model.load_state_dict(torch.load(model_path))
+    model.cuda()
+
+    # Convert to ONNX
+    convert_to_onnx(model, (1, 3, 256, 256), onnx_path)
+
+    # Build TRT engine in Python
+    convert_to_tensorrt_py(onnx_path, engine_path)
+
+
+if __name__ == '__main__':
+    main()
+