Added a more finalized model prediction

app.py

@@ -211,13 +211,20 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
     # Have to convert to image first
     result = Image.fromarray(transformed)
 
-    result_bias = str(learn_bias.predict(result))
-    result_fresh = str(learn_fresh.predict(result))
+    categories = ['Distorted', 'Maze']
+
+    def clean_output(result_values):
+        pred, idx, probs = result_values[0], result_values[1], result_values[2]
+        return dict(zip(categories, map(float, probs)))
+
+    result_bias = learn_bias.predict(result)
+    result_fresh = learn_fresh.predict(result)
     print("Results")
-    print(result_bias)
-    print(result_fresh)
+    result_bias_final = clean_output(result_bias)
+    result_fresh_final = clean_output(result_fresh)
 
-    return transformed, result_bias, result_fresh, vector_field
+    # return transformed, result_bias, result_fresh, vector_field
+    return transformed, result_bias_final, result_fresh_final, vector_field
 
 demo = gr.Interface(
     fn=transform_image,
@@ -233,15 +240,22 @@ demo = gr.Interface(
         # gr.Slider(0, 0.5, value=0.1, label="Center Smoothness")
         # gr.Checkbox(label="Reverse Gradient Direction"),
     ],
+    examples=[
+        [np.asarray(Image.open("examples/1500_maze.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
+        [np.asarray(Image.open("examples/2048_maze.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
+        [np.asarray(Image.open("examples/2300_fresh.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
+        [np.asarray(Image.open("examples/50_fresh.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5]
+    ],
     outputs=[
         gr.Image(label="Transformed Image"),
         # gr.Image(label="Result", num_top_classes=2)
-        gr.Textbox(label='Result Bias'),
-        gr.Textbox(label='Result Fresh'),
+        gr.Label(),
+        gr.Label(),
         gr.Image(label="Gradient Vector Field")
     ],
     title="Image Transformation Demo!",
-    description="This is the baseline function that will be used to generate the database for a machine learning model I am working on called 'DistortionMl'! The goal of this model is to detect and then reverse image transformations that can be generated here! You can read more about the project at this repository link : https://github.com/nick-leland/DistortionML. The main function that I was working on is the 'Bulge' function, I can't really guarantee that the others work well (;"
+    article="If you like this demo, please star the github repository for the project! Located [here!](https://github.com/nick-leland/DistortionML)",
+    description="This is the baseline function that will be used to generate the database for a machine learning model I am working on called 'DistortionMl'! The goal of this model is to detect and then reverse image transformations that can be generated here!\nYou can read more about the project at [this repository link](https://github.com/nick-leland/DistortionML). The main function that I was working on is the 'Bulge'/'Volcano' function, I can't really guarantee that the others work as well!\nI have just added the first baseline ML model to detect if a distortion has taken place! It was only trained on mazes though ([Dataset Here](https://www.kaggle.com/datasets/nickleland/distorted-mazes)) so in order for it to detect a distortion you have to use one of the images provided in the examples! Feel free to mess around wtih other images in the meantime though!"
 )
 
 demo.launch(share=True)