def test_box2d_top_left(spark: SparkSession): df = spark.createDataFrame( [ Row(values=[1.0, 2.0, 1.0, 1.0]), Row(values=[10.0, 12.0, 1.0, 5.0]), ], ["values"], ).withColumn("bbox", box2d_from_center("values")) df = df.withColumn("area", area(col("bbox"))) assert_area_equals([1.0, 5.0], df)
def test_box2d_center(spark): df = spark.createDataFrame( [ Row(values=[1.5, 2.5, 1.0, 1.0]), Row(values=[10.5, 14.5, 1.0, 5.0]), ], ["values"], ).withColumn("bbox", box2d_from_center("values")) df = df.withColumn("area", area(col("bbox"))) assert_area_equals([1.0, 5.0], df)
def test_box2d_udfs(spark): df = spark.createDataFrame( [ Row(values=[1.0, 2.0, 2.0, 3.0]), Row(values=[10.0, 12.0, 11.0, 17.0]), ], ["values"], ).withColumn("bbox", box2d("values")) df = df.withColumn("area", area(col("bbox"))) assert_area_equals([1.0, 5.0], df)
def test_areas(spark: SparkSession): """Test calculating bounding box's area.""" df = spark.createDataFrame( [ (Box2d(1, 2, 2.0, 3.0), ), (Box2d(10, 12, 11.0, 17.0), ), ], ["bbox"], ) df = df.withColumn("area", area(col("bbox"))) assert_area_equals([1.0, 5.0], df)
def test_areas(self): """Test calculating bounding box's area.""" df = self.spark.createDataFrame( [ (Box2d(1, 2, 1.0, 1.0), ), (Box2d(10, 12, 1.0, 5.0), ), ], ["bbox"], ) df = df.withColumn("area", area(col("bbox"))) self.assertCountEqual((1.0, 5.0), df.select("area").toPandas()["area"])