Demo added

Author: ign-saurav

models/experimental/panoptic_deeplab/README.md

@@ -24,6 +24,9 @@ pytest models/experimental/panoptic_deeplab/tests/test_panoptic_deeplab.py
   ```
 
 ### Demo
+```
+python models/experimental/panoptic_deeplab/demo/panoptic_deeplab_demo.py --input <input image path> --output <output image to be stored path>
+```
 **Note:** Output images will be saved in the `panoptic_deeplab_predictions/` folder.
 
 #### Single Device (BS=1):

models/experimental/panoptic_deeplab/common.py

@@ -8,6 +8,11 @@
 import pickle
 import numpy as np
 import os
+from PIL import Image
+from typing import Tuple
+import torchvision.transforms as transforms
+from typing import Optional, Any
+import ttnn
 from models.experimental.panoptic_deeplab.reference.resnet52_backbone import ResNet52BackBone as TorchBackbone
 from models.experimental.panoptic_deeplab.reference.resnet52_stem import DeepLabStem
 from torchvision.models.resnet import Bottleneck
@@ -190,12 +195,12 @@ def load_torch_model_state(torch_model: torch.nn.Module = None, layer_name: str
         model_path = model_location_generator("vision-models/panoptic_deeplab", model_subdir="", download_if_ci_v2=True)
     if model_path == "models":
         if not os.path.exists(
-            "models/experimental/panoptic_deeplab/reference/Panoptic_Deeplab_R52.pkl"
+            "models/experimental/panoptic_deeplab/resources/Panoptic_Deeplab_R52.pkl"
         ):  # check if Panoptic_Deeplab_R52.pkl is available
             os.system(
-                "models/experimental/panoptic_deeplab/reference/panoptic_deeplab_weights_download.sh"
+                "models/experimental/panoptic_deeplab/resources/panoptic_deeplab_weights_download.sh"
             )  # execute the panoptic_deeplab_weights_download.sh file
-        weights_path = "models/experimental/panoptic_deeplab/reference/Panoptic_Deeplab_R52.pkl"
+        weights_path = "models/experimental/panoptic_deeplab/resources/Panoptic_Deeplab_R52.pkl"
     else:
         weights_path = os.path.join(model_path, "Panoptic_Deeplab_R52.pkl")
 
@@ -209,7 +214,6 @@ def load_torch_model_state(torch_model: torch.nn.Module = None, layer_name: str
         if isinstance(v, np.ndarray) or isinstance(v, np.array):
             state_dict[k] = torch.from_numpy(v)
             converted_count += 1
-    logger.debug(f"Converted {converted_count} numpy arrays to torch tensors")
 
     # Get keys
     checkpoint_keys = set(state_dict.keys())
@@ -225,6 +229,9 @@ def load_torch_model_state(torch_model: torch.nn.Module = None, layer_name: str
     mapped_state_dict = {}
     for checkpoint_key, model_key in key_mapping.items():
         mapped_state_dict[model_key] = state_dict[checkpoint_key]
+    del mapped_state_dict["pixel_mean"]
+    del mapped_state_dict["pixel_std"]
+    logger.debug(f"Mapped {len(mapped_state_dict)} weights")
 
     if isinstance(
         torch_model,
@@ -240,10 +247,162 @@ def load_torch_model_state(torch_model: torch.nn.Module = None, layer_name: str
     ):
         torch_model = load_partial_state(torch_model, mapped_state_dict, layer_name)
     elif isinstance(torch_model, TorchPanopticDeepLab):
-        del mapped_state_dict["pixel_mean"]
-        del mapped_state_dict["pixel_std"]
         torch_model.load_state_dict(mapped_state_dict, strict=True)
     else:
         raise NotImplementedError("Unknown torch model. Weight loading not implemented")
 
     return torch_model.eval()
+
+
+def parameter_conv_args(torch_model: torch.nn.Module = None, parameters: dict = None):
+    from ttnn.model_preprocessing import infer_ttnn_module_args
+
+    if isinstance(torch_model, TorchPanopticDeepLab):
+        parameters.conv_args = {}
+        sample_x = torch.randn(1, 2048, 32, 64)
+        sample_res3 = torch.randn(1, 512, 64, 128)
+        sample_res2 = torch.randn(1, 256, 128, 256)
+
+        # For semantic decoder
+        if hasattr(parameters, "semantic_decoder"):
+            # ASPP
+            aspp_args = infer_ttnn_module_args(
+                model=torch_model.semantic_decoder.aspp, run_model=lambda model: model(sample_x), device=None
+            )
+            if hasattr(parameters.semantic_decoder, "aspp"):
+                parameters.semantic_decoder.aspp.conv_args = aspp_args
+
+            # Res3
+            aspp_out = torch_model.semantic_decoder.aspp(sample_x)
+            res3_args = infer_ttnn_module_args(
+                model=torch_model.semantic_decoder.res3,
+                run_model=lambda model: model(aspp_out, sample_res3),
+                device=None,
+            )
+            if hasattr(parameters.semantic_decoder, "res3"):
+                parameters.semantic_decoder.res3.conv_args = res3_args
+
+            # Res2
+            res3_out = torch_model.semantic_decoder.res3(aspp_out, sample_res3)
+            res2_args = infer_ttnn_module_args(
+                model=torch_model.semantic_decoder.res2,
+                run_model=lambda model: model(res3_out, sample_res2),
+                device=None,
+            )
+            if hasattr(parameters.semantic_decoder, "res2"):
+                parameters.semantic_decoder.res2.conv_args = res2_args
+
+            # Head
+            res2_out = torch_model.semantic_decoder.res2(res3_out, sample_res2)
+            head_args = infer_ttnn_module_args(
+                model=torch_model.semantic_decoder.head_1, run_model=lambda model: model(res2_out), device=None
+            )
+            if hasattr(parameters.semantic_decoder, "head_1"):
+                parameters.semantic_decoder.head_1.conv_args = head_args
+
+        # For instance decoder
+        if hasattr(parameters, "instance_decoder"):
+            # ASPP
+            aspp_args = infer_ttnn_module_args(
+                model=torch_model.instance_decoder.aspp, run_model=lambda model: model(sample_x), device=None
+            )
+            if hasattr(parameters.instance_decoder, "aspp"):
+                parameters.instance_decoder.aspp.conv_args = aspp_args
+
+            # Res3
+            aspp_out = torch_model.instance_decoder.aspp(sample_x)
+            res3_args = infer_ttnn_module_args(
+                model=torch_model.instance_decoder.res3,
+                run_model=lambda model: model(aspp_out, sample_res3),
+                device=None,
+            )
+            if hasattr(parameters.instance_decoder, "res3"):
+                parameters.instance_decoder.res3.conv_args = res3_args
+
+            # Res2
+            res3_out = torch_model.instance_decoder.res3(aspp_out, sample_res3)
+            res2_args = infer_ttnn_module_args(
+                model=torch_model.instance_decoder.res2,
+                run_model=lambda model: model(res3_out, sample_res2),
+                device=None,
+            )
+            if hasattr(parameters.instance_decoder, "res2"):
+                parameters.instance_decoder.res2.conv_args = res2_args
+
+            # Head
+            res2_out = torch_model.instance_decoder.res2(res3_out, sample_res2)
+            head_args_1 = infer_ttnn_module_args(
+                model=torch_model.instance_decoder.head_1, run_model=lambda model: model(res2_out), device=None
+            )
+            head_args_2 = infer_ttnn_module_args(
+                model=torch_model.instance_decoder.head_2, run_model=lambda model: model(res2_out), device=None
+            )
+            if hasattr(parameters.instance_decoder, "head_1"):
+                parameters.instance_decoder.head_1.conv_args = head_args_1
+            if hasattr(parameters.instance_decoder, "head_2"):
+                parameters.instance_decoder.head_2.conv_args = head_args_2
+    else:
+        raise NotImplementedError("Unknown torch model. Parameter conv args not implemented")
+    return parameters
+
+
+def preprocess_image(
+    image_path: str, input_width: int, input_height: int, ttnn_device: ttnn.Device, inputs_mesh_mapper: Optional[Any]
+) -> Tuple[torch.Tensor, ttnn.Tensor, np.ndarray, Tuple[int, int]]:
+    """Preprocess image for both PyTorch and TTNN"""
+    # Load image
+    image = Image.open(image_path).convert("RGB")
+    original_size = image.size  # (width, height)
+    original_array = np.array(image)
+    preprocess = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
+    )
+
+    # Resize to model input size
+    target_size = (input_width, input_height)  # PIL expects (width, height)
+    image_resized = image.resize(target_size)
+
+    # PyTorch preprocessing
+    torch_tensor = preprocess(image_resized).unsqueeze(0)  # Add batch dimension
+    torch_tensor = torch_tensor.to(torch.float)
+
+    # TTNN preprocessing
+    ttnn_tensor = None
+    ttnn_tensor = ttnn.from_torch(
+        torch_tensor.permute(0, 2, 3, 1),  # BCHW -> BHWC
+        dtype=ttnn.bfloat16,
+        device=ttnn_device,
+        mesh_mapper=inputs_mesh_mapper,
+    )
+
+    if ttnn_tensor is not None:
+        ttnn_as_torch = ttnn.to_torch(ttnn_tensor)
+
+    return torch_tensor, ttnn_tensor, original_array, original_size
+
+
+def save_preprocessed_inputs(torch_input: torch.Tensor, save_dir: str, filename: str):
+    """Save preprocessed inputs for testing purposes"""
+
+    # Create directory for test inputs
+    test_inputs_dir = os.path.join(save_dir, "test_inputs")
+    os.makedirs(test_inputs_dir, exist_ok=True)
+
+    # Save torch input tensor
+    torch_input_path = os.path.join(test_inputs_dir, f"{filename}_torch_input.pt")
+    torch.save(
+        {
+            "tensor": torch_input,
+            "shape": torch_input.shape,
+            "dtype": torch_input.dtype,
+            "mean": torch_input.mean().item(),
+            "std": torch_input.std().item(),
+            "min": torch_input.min().item(),
+            "max": torch_input.max().item(),
+        },
+        torch_input_path,
+    )
+
+    logger.info(f"Saved preprocessed torch input to: {torch_input_path}")
+
+    return torch_input_path

models/experimental/panoptic_deeplab/demo/config.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+from typing import List, Optional
+import numpy as np
+
+
+@dataclass
+class DemoConfig:
+    """Configuration class for demo parameters"""
+
+    # Model configuration
+    model_type: str = "PanopticDeepLab"
+    backbone: str = "ResNet-52"
+    num_classes: int = 19
+    weights_path: Optional[str] = None
+
+    # Input configuration
+    input_height: int = 512
+    input_width: int = 1024
+    crop_enabled: bool = False
+    normalize_enabled: bool = True
+    mean: List[float] = None
+    std: List[float] = None
+
+    # Inference configuration
+    center_threshold: float = 0.1
+    nms_kernel: int = 7
+    top_k_instances: int = 200
+    stuff_area_threshold: int = 4096
+
+    # Device configuration
+    device_id: int = 0
+    math_fidelity: str = "LoFi"
+    weights_dtype: str = "bfloat8_b"
+    activations_dtype: str = "bfloat8_b"
+
+    # Output configuration
+    save_semantic: bool = True
+    save_instance: bool = True
+    save_panoptic: bool = True
+    save_visualization: bool = True
+    save_comparison: bool = True
+
+    # Pipeline configuration
+    compare_outputs: bool = True
+    pcc_threshold: float = 0.97
+
+    # Dataset configuration (Cityscapes default)
+    thing_classes: List[int] = None
+    stuff_classes: List[int] = None
+    class_names: List[str] = None
+
+    def __post_init__(self):
+        """Initialize default values after dataclass creation"""
+        if self.mean is None:
+            self.mean = [0.485, 0.456, 0.406]
+        if self.std is None:
+            self.std = [0.229, 0.224, 0.225]
+        if self.thing_classes is None:
+            self.thing_classes = [11, 12, 13, 14, 15, 16, 17, 18]  # Cityscapes things
+        if self.stuff_classes is None:
+            self.stuff_classes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]  # Cityscapes stuff
+        if self.class_names is None:
+            self.class_names = [
+                "road",
+                "sidewalk",
+                "building",
+                "wall",
+                "fence",
+                "pole",
+                "traffic_light",
+                "traffic_sign",
+                "vegetation",
+                "terrain",
+                "sky",
+                "person",
+                "rider",
+                "car",
+                "truck",
+                "bus",
+                "train",
+                "motorcycle",
+                "bicycle",
+            ]
+
+    def _get_cityscapes_colors(self) -> np.ndarray:
+        """Get Cityscapes color palette"""
+        return np.array(
+            [
+                [128, 64, 128],  # road
+                [244, 35, 232],  # sidewalk
+                [70, 70, 70],  # building
+                [102, 102, 156],  # wall
+                [190, 153, 153],  # fence
+                [153, 153, 153],  # pole
+                [250, 170, 30],  # traffic light
+                [220, 220, 0],  # traffic sign
+                [107, 142, 35],  # vegetation
+                [152, 251, 152],  # terrain
+                [70, 130, 180],  # sky
+                [220, 20, 60],  # person
+                [255, 0, 0],  # rider
+                [0, 0, 142],  # car
+                [0, 0, 70],  # truck
+                [0, 60, 100],  # bus
+                [0, 80, 100],  # train
+                [0, 0, 230],  # motorcycle
+                [119, 11, 32],  # bicycle
+            ]
+        )