added code for pipeline-1

ASK-03 · ASK-03 · commit 52a0ec48f6cf · 2023-07-04T16:16:42.000+05:30
diff --git a/pipeline-1/capture_img.py b/pipeline-1/capture_img.py
@@ -0,0 +1,112 @@
+"""
+    This is a python script to capture images from the camera
+    and save them to the local directory in an interval of every 5 seconds.
+"""
+
+import cv2
+import time
+from datetime import datetime
+import os
+import logging
+
+
+def create_dir(**kwargs) -> None:
+    """
+    This function creates a new directory if it does not exist
+    the name of directory is images
+    """
+    import os
+
+    if kwargs is not None and "path" in kwargs:
+        path = kwargs["path"]
+    else:
+        path = "images"
+
+    format = "%(asctime)s [%(levelname)s] %(message)s"
+    logging.basicConfig(
+        level=logging.INFO,
+        format=format,
+        handlers=[logging.FileHandler("pipeline-1-capture-img.log")],
+    )
+
+    if not os.path.exists(path):
+        try:
+            os.mkdir(path)
+            logging.log(level=logging.INFO, msg=f"Directory {path} Created")
+        except Exception as e:
+            logging.log(level=logging.ERROR, msg=e)
+            logging.log(level=logging.ERROR, msg=f"Unable to create directory {path}")
+    else:
+        logging.log(level=logging.INFO, msg=f"{path} directory already exists")
+
+
+def capture_img(**kwargs) -> None:
+    """
+    This function captures images from the camera and saves them to the local directory
+    in an interval of every 5 seconds.
+    """
+
+    cap = cv2.VideoCapture(0)
+
+    if not cap.isOpened():
+        logging.log(level=logging.ERROR, msg="Unable to read camera feed")
+
+    # Default resolutions of the frame are obtained.The default resolutions are system dependent.
+    # We convert the resolutions from float to integer.
+    frame_width = int(cap.get(3))
+    frame_height = int(cap.get(4))
+
+    if not os.path.exists(os.path.join(os.getcwd(), "images")):
+        logging.log(
+            level=logging.FATAL,
+            msg="images directory not found, restart the script to create the directory",
+        )
+        exit(1)
+    else:
+        logging.log(level=logging.INFO, msg="images directory found")
+
+    while True:
+        ret, frame = cap.read()
+
+        if ret == True:
+            timestamp = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
+            filename = f"{os.getcwd()}/images/img_{timestamp}.jpg"
+
+            cv2.imwrite(filename, frame)
+            time.sleep(0.2)
+
+            if kwargs is not None and "queue" in kwargs:
+                queue = kwargs["queue"]
+                queue.put(filename)
+
+            logging.log(
+                level=logging.INFO,
+                msg=f"Image saved named img_{timestamp}.jpg in the directory images",
+            )
+
+            # Press Q on keyboard to stop recording
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                break
+
+            time.sleep(5)
+
+        else:
+            break
+
+    cap.release()
+    cv2.destroyAllWindows()
+
+
+if __name__ == "__main__":
+    format = "%(asctime)s [%(levelname)s] %(message)s"
+    logging.basicConfig(
+        level=logging.INFO,
+        format=format,
+        handlers=[logging.FileHandler("pipeline-1.log")],
+    )
+
+    logging.log(level=logging.INFO, msg="Starting the pipeline")
+    pipeline = (create_dir, capture_img)
+
+    for func in pipeline:
+        func()
diff --git a/pipeline-1/pipeline-1-main.py b/pipeline-1/pipeline-1-main.py
@@ -0,0 +1,132 @@
+# Imports
+import capture_img
+from segmentation2polygon import segmentation2polygon, do_perspective_transformation
+import queue
+import logging
+import threading
+import cv2
+from time import perf_counter, sleep
+import os
+
+"""
+TODO:
+    - we have to add two options 
+        [] - if only one image is given 
+        [done] - if we have to do it using feed from webcam
+    - the pipeline gets stopped if we get more number of points after segmentation add that.
+
+"""
+
+
+def initialize() -> None:
+    """
+    This function initializes the constant factors of the pipeline like the yolo model, logging, and to create a directory to store the images
+    """
+    format = "%(asctime)s [%(levelname)s] %(message)s"
+    logging.basicConfig(
+        level=logging.DEBUG,
+        format=format,
+        handlers=[logging.FileHandler("pipeline-1.log")],
+    )
+
+    from ultralytics import YOLO
+
+    logging.log(logging.INFO, "Loading the YOLO model")
+    global model
+    model = YOLO("./model_yolo/best.pt")
+    logging.log(logging.INFO, "Loaded the YOLO model")
+
+    logging.log(level=logging.INFO, msg="Creating the directory structure!")
+    if not os.path.exists("images"):
+        capture_img.create_dir(path="images")
+    if not os.path.exists("processed_images"):
+        capture_img.create_dir(path="processed_images")
+    logging.log(level=logging.DEBUG, msg=f"{os.getcwd()}")
+    logging.log(level=logging.INFO, msg="Created the directory structure!")
+
+    sleep(2)
+
+
+def get_segmentation(img_path: str):
+    """
+    This function returns the segmentation mask of the image
+    Input: img_path: str
+    Output: segmentation: np.ndarray
+    """
+    logging.log(logging.INFO, "Getting the segmentation mask of the image")
+    start = perf_counter()
+    results = model.predict(img_path)
+    end = perf_counter()
+    logging.log(
+        logging.INFO,
+        f"Time taken to get the segmentation mask of the image: {end-start}",
+    )
+
+    try:
+        mask = results[0].masks
+        mask = mask.xy
+        segmentation = segmentation2polygon(mask[0])
+        print(segmentation.reshape((-1, 2)))
+        return segmentation.reshape((-1, 2))
+    except Exception as e:
+        logging.log(logging.ERROR, e)
+        return None
+
+
+def main() -> None:
+    img_queue = queue.Queue()
+
+    img_capture_thread = threading.Thread(
+        target=capture_img.capture_img, kwargs={"queue": img_queue}, daemon=True
+    )
+    img_capture_thread.start()
+
+    while True:
+        if 0xFF == ord("q"):
+            break
+
+        if not img_queue.empty():
+            img_path = img_queue.get()
+            logging.log(level=logging.DEBUG, msg=f"{img_queue.qsize()}")
+            logging.log(level=logging.DEBUG, msg=f"Got the image path: {img_path}")
+            segmentation = get_segmentation(img_path=img_path)
+            logging.log(
+                level=logging.INFO, msg="Got the segmentation mask of the image"
+            )
+
+            if segmentation is not None:
+                img = cv2.imread(filename=img_path)
+                processed_image = do_perspective_transformation(
+                    image=img, input_array=segmentation
+                )
+                logging.log(level=logging.INFO, msg="Got the processed image")
+
+                try:
+                    if not os.path.exists("processed_images"):
+                        logging.log(
+                            level=logging.INFO,
+                            msg="Creating the directory to store the processed images",
+                        )
+                        os.mkdir("processed_images")
+                    img_name = img_path.split("/")[-1]
+                    cv2.imwrite(
+                        filename=f"processed_images/{img_name}", img=processed_image
+                    )
+                    logging.log(level=logging.INFO, msg="Saved the processed image")
+
+                except Exception as e:
+                    logging.log(level=logging.ERROR, msg=e)
+                    logging.log(
+                        level=logging.ERROR, msg="Unable to save the processed image"
+                    )
+            img_queue.task_done()
+
+    img_capture_thread.join()
+    return
+
+
+if __name__ == "__main__":
+    initialize()
+    main()
+    del model
+    exit(0)
diff --git a/pipeline-1/segmentation2polygon.py b/pipeline-1/segmentation2polygon.py
@@ -0,0 +1,143 @@
+# Imports
+import numpy as np
+import cv2
+
+
+# Functions
+def get_segmentation():
+    """
+    HELPER FUNCTION
+    This function returns the segmentation mask of the image
+    """
+    return "0 0.55 0.255556 0.548437 0.258333 0.521875 0.258333 0.520312 0.261111 0.515625 0.261111 0.514063 0.263889 0.50625 0.263889 0.504687 0.266667 0.496875 0.266667 0.495313 0.269444 0.489062 0.269444 0.4875 0.272222 0.476562 0.272222 0.475 0.275 0.448438 0.275 0.446875 0.277778 0.432813 0.277778 0.43125 0.280556 0.43125 0.705556 0.432813 0.708333 0.50625 0.708333 0.507812 0.705556 0.51875 0.705556 0.520312 0.702778 0.529688 0.702778 0.53125 0.7 0.546875 0.7 0.548437 0.702778 0.69375 0.702778 0.695312 0.705556 0.715625 0.705556 0.717188 0.708333 0.734375 0.708333 0.735937 0.705556 0.7375 0.705556 0.7375 0.7 0.739062 0.697222 0.739062 0.680556 0.740625 0.677778 0.740625 0.616667 0.742188 0.613889 0.742188 0.258333 0.740625 0.255556"
+
+
+def convert_str_to_array(string) -> list[float]:
+    """
+    HELPER FUNCTION
+    This function converts the string to a list of floats
+    Input: string: string of points
+    Output: list of floats
+    """
+    return [float(i) for i in string.split(" ")]
+
+
+def get_height_and_width_of_img(img_path) -> dict[str, int]:
+    """
+    HELPER FUNCTION
+    This function returns the height and width of the image
+    Input: img_path: path to the image
+    Output: dict of height and width
+    """
+    img = cv2.imread(img_path)
+    h, w, c = img.shape
+    return {"height": h, "width": w}
+
+
+def from_normalised_to_pixel(img_size: dict, contour):
+    """
+    function converts the normalised segmentation mask to pixel values
+    Input: img_size: dict of height and width of image
+           contour: list of points
+    Output: res: list of points
+    """
+    height = img_size["height"]
+    width = img_size["width"]
+    res = []
+
+    for i in range(len(contour[1::2])):
+        x = int(contour[2 * i + 1] * width)
+        y = int(contour[2 * i + 2] * height)
+        res.append((x, y))
+
+    return res
+
+
+def segmentation2polygon(segmentation):
+    """
+    function converts the segmentation mask from yolo V8 to a 4 sided polygon
+    Input: segmentation: list of points
+    Output: approx: list of points
+    """
+
+    points = np.array(segmentation, dtype=np.int32)
+
+    closed = True
+
+    peri = cv2.arcLength(points, True)
+    approx = cv2.approxPolyDP(points, 0.02 * peri, closed)
+
+    return approx
+
+
+def pipeline():
+    """
+    Function is the pipeline for the segmentation to polygon conversion
+    """
+    segmentation = get_segmentation()
+    segmentation = convert_str_to_array(segmentation)
+    img_size = get_height_and_width_of_img(
+        "./yolo/yolo_predictions/medicakolkata_rcu_mon--8_2022_5_18_17_5_38.jpeg"
+    )
+    polygon = from_normalised_to_pixel(img_size, segmentation)
+    polygon = segmentation2polygon(polygon)
+
+    return polygon
+
+
+def draw_points(polygon) -> None:
+    """
+    Function to draw the points on the image, it shows the image in the new window with points
+    Input: polygon: list of points
+    Output: None
+    """
+    img = cv2.imread(
+        "./yolo/yolo_predictions/medicakolkata_rcu_mon--8_2022_5_18_17_5_38.jpeg"
+    )
+    for i in polygon:
+        cv2.circle(img, (i[0], i[1]), 5, (0, 0, 255), -1)
+
+    cv2.imshow("img", img)
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()
+
+
+def do_perspective_transformation(image, input_array):
+    """
+    Function to do perspective transformation
+    Input: image: image on which perspective transformation is to be done
+           input_array: list of points
+    Output: result: image after perspective transformation
+    """
+    height, width = image.shape[:2]
+
+    input_array = np.array(input_array, dtype=np.float32)
+    output_array = np.array(
+        [(0, 0), (0, height), (width, height), (width, 0)], dtype=np.float32
+    )
+
+    matrix = cv2.getPerspectiveTransform(input_array, output_array)
+    result = cv2.warpPerspective(
+        image,
+        matrix,
+        (width, height),
+        borderMode=cv2.BORDER_CONSTANT,
+        borderValue=(0, 0, 0),
+    )
+
+    return result
+
+
+if __name__ == "__main__":
+    approx_polygon = pipeline().reshape((-1, 2))
+    print(len(approx_polygon))
+    draw_points(approx_polygon)
+    perspective = do_perspective_transformation(
+        cv2.imread(
+            "./yolo/yolo_predictions/medicakolkata_rcu_mon--8_2022_5_18_17_5_38.jpeg"
+        ),
+        approx_polygon,
+    )
+    cv2.imshow("perspective", perspective)
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()
