Update src/utils.py

src/utils.py

@@ -1,90 +1,92 @@
-import numpy as np
-import cv2
-from PIL import Image
-
-def convert_to_bw(image):
-    """
-    Converts a PIL image to black & white (grayscale),
-    and then back to RGB to maintain compatibility with other processes.
-    
-    Parameters:
-        image (PIL.Image): Input RGB image.
-
-    Returns:
-        PIL.Image: Black & white image in RGB format.
-    """
-    return image.convert("L").convert("RGB")
-
-def load_colorization_model():
-    """
-    Loads the pre-trained Caffe model for colorizing black & white images.
-    
-    Model files required:
-        - colorization_deploy_v2.prototxt
-        - colorization_release_v2.caffemodel
-        - pts_in_hull.npy
-
-    Returns:
-        cv2.dnn_Net: Loaded and initialized OpenCV DNN colorization model.
-    """
-    # Paths to model architecture, weights, and cluster centers
-    proto_file = "models/colorization_deploy_v2.prototxt"
-    model_file = "models/colorization_release_v2.caffemodel"
-    cluster_file = "models/pts_in_hull.npy"
-
-
-    # Load the model using OpenCV DNN module
-    net = cv2.dnn.readNetFromCaffe(proto_file, model_file)
-    pts = np.load(cluster_file)
-
-    # Populate cluster centers as 1x1 convolution kernel
-    class8_ab = net.getLayerId("class8_ab")
-    conv8_313_rh = net.getLayerId("conv8_313_rh")
-
-    pts = pts.transpose().reshape(2, 313, 1, 1)
-    net.getLayer(class8_ab).blobs = [pts.astype(np.float32)]
-    net.getLayer(conv8_313_rh).blobs = [np.full([1, 313], 2.606, dtype=np.float32)]
-
-    return net
-
-def colorize_bw_image(pil_img, net):
-    """
-    Colorizes a grayscale (black & white) image using a pre-trained DNN model.
-
-    Parameters:
-        pil_img (PIL.Image): Input grayscale image in RGB format.
-        net (cv2.dnn_Net): Loaded OpenCV DNN colorization model.
-
-    Returns:
-        PIL.Image: Colorized image in RGB format.
-    """
-    # Convert PIL image to NumPy array
-    img = np.array(pil_img)
-    img_rgb = img[:, :, [2, 1, 0]]  # Convert RGB to BGR
-    img_rgb = img_rgb.astype("float32") / 255.0
-    # Convert to LAB color space and extract L channel
-    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2LAB)
-    l_channel = img_lab[:, :, 0]
-    
-    # Resize L channel to match model input size and normalize
-    input_l = cv2.resize(l_channel, (224, 224))
-    input_l -= 50
-    
-    # Run inference
-    net.setInput(cv2.dnn.blobFromImage(input_l))
-    ab_channels = net.forward()[0, :, :, :].transpose((1, 2, 0))  # shape: (56, 56, 2)
-
-    # Resize predicted ab channels to match original image size
-    ab_channels = cv2.resize(ab_channels, (img.shape[1], img.shape[0]))
-    
-    # Merge original L channel with predicted ab channels
-    lab_output = np.concatenate((l_channel[:, :, np.newaxis], ab_channels), axis=2)
-    
-    # Convert LAB to BGR, clip values, and convert to uint8
-    bgr_out = cv2.cvtColor(lab_output, cv2.COLOR_LAB2BGR)
-    bgr_out = np.clip(bgr_out, 0, 1)
-
-    
-    # Convert back to RGB and return as PIL Image
-    final_rgb = (bgr_out[:, :, [2, 1, 0]] * 255).astype("uint8")
-    return Image.fromarray(final_rgb)
+import numpy as np
+import cv2
+from PIL import Image
+import os
+
+def convert_to_bw(image):
+    """
+    Converts a PIL image to black & white (grayscale),
+    and then back to RGB to maintain compatibility with other processes.
+    
+    Parameters:
+        image (PIL.Image): Input RGB image.
+
+    Returns:
+        PIL.Image: Black & white image in RGB format.
+    """
+    return image.convert("L").convert("RGB")
+
+def load_colorization_model():
+    """
+    Loads the pre-trained Caffe model for colorizing black & white images.
+    
+    Model files required:
+        - colorization_deploy_v2.prototxt
+        - colorization_release_v2.caffemodel
+        - pts_in_hull.npy
+
+    Returns:
+        cv2.dnn_Net: Loaded and initialized OpenCV DNN colorization model.
+    """
+    # Paths to model architecture, weights, and cluster centers
+    base_path = os.path.join(os.path.dirname(__file__), "models")
+    proto_file = "models/colorization_deploy_v2.prototxt"
+    model_file = "models/colorization_release_v2.caffemodel"
+    cluster_file = "models/pts_in_hull.npy"
+
+
+    # Load the model using OpenCV DNN module
+    net = cv2.dnn.readNetFromCaffe(proto_file, model_file)
+    pts = np.load(cluster_file)
+
+    # Populate cluster centers as 1x1 convolution kernel
+    class8_ab = net.getLayerId("class8_ab")
+    conv8_313_rh = net.getLayerId("conv8_313_rh")
+
+    pts = pts.transpose().reshape(2, 313, 1, 1)
+    net.getLayer(class8_ab).blobs = [pts.astype(np.float32)]
+    net.getLayer(conv8_313_rh).blobs = [np.full([1, 313], 2.606, dtype=np.float32)]
+
+    return net
+
+def colorize_bw_image(pil_img, net):
+    """
+    Colorizes a grayscale (black & white) image using a pre-trained DNN model.
+
+    Parameters:
+        pil_img (PIL.Image): Input grayscale image in RGB format.
+        net (cv2.dnn_Net): Loaded OpenCV DNN colorization model.
+
+    Returns:
+        PIL.Image: Colorized image in RGB format.
+    """
+    # Convert PIL image to NumPy array
+    img = np.array(pil_img)
+    img_rgb = img[:, :, [2, 1, 0]]  # Convert RGB to BGR
+    img_rgb = img_rgb.astype("float32") / 255.0
+    # Convert to LAB color space and extract L channel
+    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2LAB)
+    l_channel = img_lab[:, :, 0]
+    
+    # Resize L channel to match model input size and normalize
+    input_l = cv2.resize(l_channel, (224, 224))
+    input_l -= 50
+    
+    # Run inference
+    net.setInput(cv2.dnn.blobFromImage(input_l))
+    ab_channels = net.forward()[0, :, :, :].transpose((1, 2, 0))  # shape: (56, 56, 2)
+
+    # Resize predicted ab channels to match original image size
+    ab_channels = cv2.resize(ab_channels, (img.shape[1], img.shape[0]))
+    
+    # Merge original L channel with predicted ab channels
+    lab_output = np.concatenate((l_channel[:, :, np.newaxis], ab_channels), axis=2)
+    
+    # Convert LAB to BGR, clip values, and convert to uint8
+    bgr_out = cv2.cvtColor(lab_output, cv2.COLOR_LAB2BGR)
+    bgr_out = np.clip(bgr_out, 0, 1)
+
+    
+    # Convert back to RGB and return as PIL Image
+    final_rgb = (bgr_out[:, :, [2, 1, 0]] * 255).astype("uint8")
+    return Image.fromarray(final_rgb)