def initialize()

in community-content/vertex_model_garden/model_oss/vot/handler.py [0:0]
34 lines of code
8 McCabe index (conditional complexity)

  def initialize(self, context: Any) -> None:
    properties = context.system_properties
    self.map_location = (
        "cuda"
        if torch.cuda.is_available() and properties.get("gpu_id") is not None
        else "cpu"
    )
    self.device = torch.device(
        self.map_location + ":" + str(properties.get("gpu_id"))
        if torch.cuda.is_available() and properties.get("gpu_id") is not None
        else self.map_location
    )
    self.manifest = context.manifest

    detection_endpoint_id = os.environ.get("DETECTION_ENDPOINT", None)
    if detection_endpoint_id:
      self.detection_endpoint = aiplatform.Endpoint(detection_endpoint_id)
      endpoint_label_map = os.environ.get("LABEL_MAP", None)
      if endpoint_label_map:
        endpoint_label_map_file = endpoint_label_map
        self.label_map = commons.get_label_map(endpoint_label_map_file)
      else:
        raise ValueError(
            "LABEL MAP must be provided with DETECTION ENDPOINT:"
            f" {self.detection_endpoint}"
        )

    self.track_thresh = os.environ.get("TRACK_THRESHOLD", _TRACK_THRESHOLD)
    self.track_buffer = os.environ.get("TRACK_BUFFER", _TRACK_BUFFER)
    self.match_thresh = os.environ.get("MATCH_THRESHOLD", _MATCH_THRESHOLD)
    self.save_video_results = bool(int(os.environ.get("SAVE_VIDEO_RESULTS", 0)))
    self.output_bucket = os.environ.get("OUTPUT_BUCKET", None)
    if not self.output_bucket:
      raise ValueError("Empty Output Bucket.")
    self.initialized = True
    logging.info("Handler initialization done.")