fix: normalize image for output

app.py

@@ -47,10 +47,12 @@ def segment_image(image):
     image = PILImage.create(image)
     prediction, _, _ = LEARNER.predict(image)
 
-    print("Prediction shape:", prediction.shape)
-    print("Unique values:", numpy.unique(prediction))
+    prediction_array = numpy.array(prediction, dtype=numpy.uint8)
 
-    return numpy.array(prediction, dtype=numpy.uint8)
+    # Normalize class indices to 0-255 for proper visualization
+    prediction_array = (prediction_array / prediction_array.max()) * 255
+
+    return prediction_array.astype(numpy.uint8)
 
 
 demo = gradio.Interface(