quick fix

app.py

@@ -4,19 +4,31 @@ import numpy as np
 import glob
 import warnings
 import pandas as pd
-from .utils import OrthogonalRegularizer
+import matplotlib.pyplot as plt
 
+from utils import OrthogonalRegularizer
 from huggingface_hub.keras_mixin import from_pretrained_keras
 
 # load model
-model = from_pretrained_keras("keras-io/pointnet_segmentation",  custom_objects={"OrthogonalRegularizer": OrthogonalRegularizer})
+model = from_pretrained_keras(
+    "keras-io/pointnet_segmentation", custom_objects={"OrthogonalRegularizer": OrthogonalRegularizer}
+)
 
 # Examples
 samples = []
-input_images = glob.glob('asset/source/*.png')
-examples = [[im, f"asset/source/{im.split('/')[-1].split('.')[0]}.csv", f'asset/ground_truth/{im_name}.png'] for im in input_images]
-LABELS = ['wing', 'body', 'tail', 'engine']
-    
+input_images = glob.glob("asset/source/*.png")
+examples = [
+    [
+        im,
+        f"asset/ground_truth/{im.split('/')[-1].split('.')[0]}.png",
+        f"asset/source/{im.split('/')[-1].split('.')[0]}.csv",
+    ]
+    for im in input_images
+]
+LABELS = ["wing", "body", "tail", "engine"]
+COLORS = ["blue", "green", "red", "pink"]
+
+
 def visualize_data(point_cloud, labels, output_path=None):
     df = pd.DataFrame(
         data={
@@ -31,9 +43,7 @@ def visualize_data(point_cloud, labels, output_path=None):
     for index, label in enumerate(LABELS):
         c_df = df[df["label"] == label]
         try:
-            ax.scatter(
-                c_df["x"], c_df["y"], c_df["z"], label=label, alpha=0.5, c=COLORS[index]
-            )
+            ax.scatter(c_df["x"], c_df["y"], c_df["z"], label=label, alpha=0.5, c=COLORS[index])
         except IndexError:
             pass
     ax.legend()
@@ -41,45 +51,45 @@ def visualize_data(point_cloud, labels, output_path=None):
         os.makedirs(os.path.dirname(output_path), exist_ok=True)
         plt.savefig(output_path)
 
-def inference(im_path,
-              truth_path,
-              file_obj,
-              output_path = 'asset/output',
-              cpu = False,
-              ):
+
+def inference(
+    im_path,
+    truth_path,
+    file_obj,
+    output_path="asset/output",
+    cpu=False,
+):
 
     csv_path = file_obj.name
-    im_name = csv_path.split('/')[-1].split('.')[0]
-    
+    im_name = csv_path.split("/")[-1].split(".")[0]
+
     if os.path.exists(csv_path):
         df = pd.read_csv(csv_path, index_col=None)
-        inputs = df[['x', 'y', 'z']].values
+        inputs = df[["x", "y", "z"]].values
         y_test = df.iloc[:, 3:].values
     else:
-        warnings.warn(f'{csv_path} not found for {im_path}')
+        warnings.warn(f"{csv_path} not found for {im_path}")
         return
 
-
     preds = model.predict(np.expand_dims(inputs, 0))[0]
     label_map = LABELS + ["none"]
-    visualize_data(inputs, [label_map[np.argmax(label)] for label in preds], f'{output_path}/{im_name}.png')
-    return f'{output_path}/{im_name}.png'
+    visualize_data(inputs, [label_map[np.argmax(label)] for label in preds], f"{output_path}/{im_name}.png")
+    return f"{output_path}/{im_name}.png"
+
 
 article = "<div style='text-align: center;'><a href='https://nouamanetazi.me/' target='_blank'>Space by Nouamane Tazi</a><br><a href='https://keras.io/examples/vision/pointnet_segmentation' target='_blank'>Keras example by Soumik Rakshit, Sayak Paul</a></div>"
 
 iface = gr.Interface(
-    inference, # main function
-    inputs = [ 
-        gr.inputs.Image(label='Image', type="filepath"),
-        gr.inputs.Image(label='Ground Truth', type="filepath"),
-        "file"
-        
+    inference,  # main function
+    inputs=[
+        gr.inputs.Image(label="Image", type="filepath"),
+        gr.inputs.Image(label="Ground Truth", type="filepath"),
+        "file",
+    ],
+    outputs=[
+        gr.outputs.Image(label="result"),  # generated image
     ],
-    outputs = [
-        gr.outputs.Image(label='result'), # generated image
-    ], 
-    
-    title = 'Point cloud segmentation with PointNet',
-    article = article,
-    examples = examples,
-).launch(enable_queue=True, cache_examples=True)
+    title="Point cloud segmentation with PointNet",
+    article=article,
+    examples=examples,
+).launch(enable_queue=True, cache_examples=True)

utils.py

@@ -1,5 +1,7 @@
+import tensorflow as tf
 from tensorflow import keras
 
+
 class OrthogonalRegularizer(keras.regularizers.Regularizer):
     """Reference: https://keras.io/examples/vision/pointnet/#build-a-model"""
 
@@ -12,9 +14,7 @@ class OrthogonalRegularizer(keras.regularizers.Regularizer):
         identity = tf.cast(self.identity, x.dtype)
         x = tf.reshape(x, (tf.shape(x)[0], self.num_features, self.num_features))
         xxt = tf.tensordot(x, x, axes=(2, 2))
-        xxt = tf.reshape(
-            xxt, (tf.shape(x)[0] * tf.shape(x)[0], self.num_features, self.num_features)
-        )
+        xxt = tf.reshape(xxt, (tf.shape(x)[0] * tf.shape(x)[0], self.num_features, self.num_features))
         return tf.reduce_sum(self.l2reg * tf.square(xxt - identity))
 
     def get_config(self):