Fix KeyError in create_geojson Function Due to Incorrect cid Handling

[neelrajpatil]

Description:

The create_geojson function in viz_utils.py throws an KeyError error because cid is always a list or tuple. The error occurs because the code attempts to directly use cid as a key to the lookup dict, leading to incorrect handling and subsequent failures.

Error Message:

Traceback (most recent call last):
  File "path/to/main.py", line 232, in <module>
    main(params)
  File "path/to/main.py", line 117, in main
    z_pp = post_process_main(
           ^^^^^^^^^^^^^^^^^^
  File "path/to/src/post_process.py", line 95, in post_process_main
    create_polygon_output(pinst_out, pcls_out, params["output_dir"], params)
  File "path/to/src/viz_utils.py", line 118, in create_polygon_output
    create_geojson(
  File "path/to/src/viz_utils.py", line 34, in create_geojson
    "name": lookup[cid],
            ~~~~~~^^^^^
KeyError: (6, (1248.144578313253, 9795.885542168675))

Suggested Fix:

Checking classids is a list/tuple and extracting the first element of cid.

Current Code:

def create_geojson(polygons, classids, lookup, result_dir):
    features = []
    for i, (poly, cid) in enumerate(zip(polygons, classids)):
        poly = np.array(poly)
        poly = poly[:, [1, 0]]
        poly = poly.tolist()
        # poly.append(poly[0])
        feature = geojson.Feature(
            geometry=geojson.LineString(poly),
            properties={
                "Name": f"Nuc {i}",
                "Type": "Polygon",
                "classification": {
                    "name": lookup[cid],
                    "color": COLORS_LIZARD[cid - 1],
                },
            },
        )
        features.append(feature)
    feature_collection = geojson.FeatureCollection(features)
    geojson_str = geojson.dumps(feature_collection, indent=2)
    with open(result_dir + "/poly.geojson", "w") as f:
        f.write(geojson_str)

Proposed Code:

def create_geojson(polygons, classids, lookup, result_dir):
    features = []
    if isinstance(classids[0], (list, tuple)):
        classids = [cid[0] for cid in classids]
    for i, (poly, cid) in enumerate(zip(polygons, classids)):
        poly = np.array(poly)
        poly = poly[:, [1, 0]]
        poly = poly.tolist()
        # poly.append(poly[0])
        feature = geojson.Feature(
            geometry=geojson.LineString(poly),
            properties={
                "Name": f"Nuc {i}",
                "Type": "Polygon",
                "classification": {
                    "name": lookup[cid],
                    "color": COLORS_LIZARD[cid - 1],
                },
            },
        )
        features.append(feature)
    feature_collection = geojson.FeatureCollection(features)
    geojson_str = geojson.dumps(feature_collection, indent=2)
    with open(result_dir + "/poly.geojson", "w") as f:
        f.write(geojson_str)```

[eliasbaumann]

Thanks for catching this, will merge!

I think the polygon/geojson output needs to be properly integrated, perhaps already at class-assignment so that this re-search of all objects is not necessary.