Updated the app to now provide a 'fixed' image

app.py

@@ -276,8 +276,8 @@ learn_fresh = load_learner('model_fresh.pkl')
 
 # Loads the YOLO Model
 model_bulge = YOLO("bulge_yolo_model.pt")
-modelv8x = YOLO("yolov8x.pt")
-modelv8n = YOLO("yolov8n.pt")
+# modelv8x = YOLO("yolov8x.pt")
+# modelv8n = YOLO("yolov8n.pt")
 
 def predict_image(img, model, conf_threshold, iou_threshold):
     """Predicts objects in an image using a YOLOv8 model with adjustable confidence and IOU thresholds."""
@@ -300,7 +300,7 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
     I = np.asarray(Image.open(image))    
 
     # Downsample large images
-    max_size = 1024  # Increased from 512 to allow for more detail
+    max_size = 640 # Increased from 512 to allow for more detail, decreased from 1024 to match YOLO model training.
     if max(I.shape[:2]) > max_size:
         scale = max_size / max(I.shape[:2])
         new_size = (int(I.shape[1] * scale), int(I.shape[0] * scale))
@@ -352,6 +352,7 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
         func = bulge
         edge_smoothness = 0
         center_smoothness = 0
+
     elif func_choice == "Volcano":
         func = bulge
         edge_smoothness = 0
@@ -388,11 +389,14 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
         inv_gx, inv_gy = -gx, -gy
         x_inv = x + inv_gx
         y_inv = y + inv_gy
+
         x_inv = np.clip(x_inv, 0, cols - 1)
         y_inv = np.clip(y_inv, 0, rows - 1)
         
         inverse_transformed = cv2.remap(I, x_inv, y_inv, cv2.INTER_LINEAR)
 
+        applied_transformed = cv2.remap(transformed, x_inv, y_inv, cv2.INTER_LINEAR)
+
         print(f"Transformed image shape: {transformed.shape}")
         print(f"Inverse transformed image shape: {inverse_transformed.shape}")
         
@@ -407,6 +411,8 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
             original_size = np.asarray(Image.open(image)).shape[:2][::-1]
             transformed = cv2.resize(transformed, original_size, interpolation=cv2.INTER_LINEAR)
             inverse_transformed = cv2.resize(inverse_transformed, original_size, interpolation=cv2.INTER_LINEAR)
+            applied_transformed = cv2.resize(applied_transformed, original_size, interpolation=cv2.INTER_LINEAR)
+
             vector_field = cv2.resize(vector_field, original_size, interpolation=cv2.INTER_LINEAR)
             inverted_vector_field = cv2.resize(inverted_vector_field, original_size, interpolation=cv2.INTER_LINEAR)
 
@@ -420,30 +426,33 @@ def transform_image(image, func_choice, randomization_check, radius, center_x, c
 
     result = Image.fromarray(transformed)
 
-    categories = ['Distorted', 'Maze']
+    # categories = ['Distorted', 'Maze']
 
-    def clean_output(result_values):
-        pred, idx, probs = result_values
-        return dict(zip(categories, map(float, probs)))
+    # def clean_output(result_values):
+    #     pred, idx, probs = result_values
+    #     return dict(zip(categories, map(float, probs)))
 
-    result_bias = learn_bias.predict(result)
-    result_fresh = learn_fresh.predict(result)
-    result_bias_final = clean_output(result_bias)
-    result_fresh_final = clean_output(result_fresh)
+    # Outdated, changing to a classification basis
+    # result_bias = learn_bias.predict(result)
+    # result_fresh = learn_fresh.predict(result)
+    # result_bias_final = clean_output(result_bias)
+    # result_fresh_final = clean_output(result_fresh)
 
     result_localization = model_bulge.predict(transformed, save=True)
     print(result_localization, "bulge")
-    result_localization1 = modelv8n.predict(transformed, save=True)
-    print(result_localization1, "modelv8n")
-    result_localization2 = modelv8x.predict(transformed, save=True)
-    print(result_localization2, "modelv8x")
+    # result_localization1 = modelv8n.predict(transformed, save=True)
+    # print(result_localization1, "modelv8n")
+    # result_localization2 = modelv8x.predict(transformed, save=True)
+    # print(result_localization2, "modelv8x")
 
 
     YOLO_image = predict_image(transformed, model_bulge, 0.5, 0.5)
-    YOLO_image1 = predict_image(transformed, modelv8n, 0.5, 0.5)
-    YOLO_image2 = predict_image(transformed, modelv8x, 0.5, 0.5)
+    # YOLO_image1 = predict_image(transformed, modelv8n, 0.5, 0.5)
+    # YOLO_image2 = predict_image(transformed, modelv8x, 0.5, 0.5)
 
-    return transformed, YOLO_image, YOLO_image1, YOLO_image2, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field
+    # return transformed, YOLO_image, YOLO_image1, YOLO_image2, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field
+    # return transformed, YOLO_image, result_bias_final, result_fresh_final, vector_field, inverse_transformed, inverted_vector_field
+    return transformed, YOLO_image, vector_field, inverse_transformed, inverted_vector_field, applied_transformed
 
 
 demo = gr.Interface(
@@ -469,13 +478,14 @@ demo = gr.Interface(
     outputs=[
         gr.Image(label="Transformed Image"),
         gr.Image(label="bulge_model Model Classification"),
-        gr.Image(label="yolov8n Model Classification"),
-        gr.Image(label="yolov8x Model Classification"),
-        gr.Label(),
-        gr.Label(),
+        # gr.Image(label="yolov8n Model Classification"),
+        # gr.Image(label="yolov8x Model Classification"),
+        # gr.Label(),
+        # gr.Label(),
         gr.Image(label="Gradient Vector Field"),
         gr.Image(label="Inverse Gradient"),
         gr.Image(label="Inverted Vector Field"),
+        gr.Image(label="Fixed Image")
     ],
     title="Image Transformation Demo!",
     article="If you like this demo, please star the github repository for the project! Located [here!](https://github.com/nick-leland/DistortionML)",