init commit

extract_data.py

@@ -0,0 +1,163 @@
+import os
+from datetime import datetime
+
+import cv2 as cv
+import numpy as np
+import pandas as pd
+import pytesseract as tess
+from PIL import Image
+
+PHOTOS_PATH = "./photos/"
+
+
+# Get a list of images given a directory path
+def get_images(url: str):
+    images = []
+    for img_url in os.listdir(url):
+        try:
+            image = Image.open(os.path.join(url, img_url))
+            images.append(image)
+        except IOError:
+            print(f"Error opening image: {img_url}")
+    return images
+
+
+# Get the datetime taken from an image
+def get_datetime_taken(image: Image.Image) -> datetime | None:
+    exif = image.getexif()
+    if 306 in exif:
+        return datetime.strptime(exif[306], "%Y:%m:%d %H:%M:%S")
+    return None
+
+
+# Convert an image to OpenCV format
+def convert_to_opencv_image(img: Image.Image) -> np.ndarray:
+    return cv.cvtColor(np.array(img), cv.COLOR_RGB2BGR)
+
+
+def order_points(pts):
+    pts = pts.reshape(4, 2)
+    rect = np.zeros((4, 2), dtype="float32")
+
+    s = pts.sum(axis=1)
+    rect[0] = pts[np.argmin(s)]  # top-left
+    rect[2] = pts[np.argmax(s)]  # bottom-right
+
+    diff = np.diff(pts, axis=1)
+    rect[1] = pts[np.argmin(diff)]  # top-right
+    rect[3] = pts[np.argmax(diff)]  # bottom-left
+
+    return rect
+
+
+def is_closed_contour(cnt, eps=1.0):
+    # Check area
+    if cv.contourArea(cnt) == 0:
+        return False
+    # Check if first and last points are close
+    return cv.norm(cnt[0][0] - cnt[-1][0]) < eps
+
+
+# Optimise the image for OCR
+def process_image(img: Image.Image):
+    arr = convert_to_opencv_image(img)
+
+    # Blur the image for better edge (contour) detection
+    blur = cv.GaussianBlur(arr, (7, 7), 0)
+    edges = cv.Canny(blur, 50, 100)
+    contours, hierarchy = cv.findContours(
+        edges, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE
+    )
+
+    # Filter contours for rectangles
+    candidates = []
+    img_area = arr.shape[0] * arr.shape[1]
+
+    for cnt in contours:
+        rect = cv.minAreaRect(cnt)
+        (center, (width, height), angle) = rect
+        box = cv.boxPoints(rect)
+        box_contour = box.reshape((-1, 1, 2))
+        area = cv.contourArea(box_contour)
+        if area < 0.01 * img_area:
+            continue
+        # Check the aspect ratio is reasonable
+        aspect_ratio = width / float(height)
+
+        if 0.9 < aspect_ratio < 1.1:
+            candidates.append(box_contour)
+
+    # Most likely rectangle will be the largest one
+    if len(candidates) == 0:
+        return None
+
+    cv.drawContours(arr, contours, -1, (0, 255, 0), 3)
+    preview_image(arr)
+
+    display_contour = max(candidates, key=cv.contourArea)
+    rect = order_points(display_contour)
+    (w, h) = (400, 400)
+    dst = np.array([[0, 0], [w - 1, 0], [w - 1, h - 1], [0, h - 1]], dtype="float32")
+    mat = cv.getPerspectiveTransform(rect, dst)
+    warped = cv.warpPerspective(arr, mat, (w, h))
+    return warped
+
+
+# Get the text from an image using OCR
+def ocr_image(img: Image.Image) -> str:
+    return None
+
+
+# Process OCR text output
+def process_ocr_text(text: str) -> str:
+    return None
+
+
+# There are two gyms that I go to, one is the Peckham gym and the other is the Elephant and Castle gym.
+# You can tell which gym I went to by looking at the color of the wall.
+# If there is a green wall, its most likely the Peckham gym.
+# If there is a blue wall, its most likely the Elephant and Castle gym.
+def get_gym(image: Image.Image) -> str | None:
+    return None
+
+
+def preview_image(img: np.ndarray):
+    cv.imshow("preview", img)
+    cv.waitKey(0)
+    cv.destroyAllWindows()
+
+
+def __main__():
+    imgs = get_images(PHOTOS_PATH)
+
+    if not imgs:
+        print("No images")
+        return None
+
+    fail = []
+    success = []
+    for img in imgs:
+        tst = process_image(img)
+        if tst is None:
+            fail.append(img.filename)
+            continue
+
+        success.append(img.filename)
+    print("success_len: ", len(success))
+    print("fail_len: ", len(fail))
+
+    print("failed:")
+    for x in fail:
+        print(x)
+
+    print("success:")
+    for x in success:
+        print(x)
+
+    print("success_len: ", len(success))
+    print("fail_len: ", len(fail))
+
+    return None
+
+
+__main__()