168 lines
4.2 KiB
Python
168 lines
4.2 KiB
Python
import os
|
|
from datetime import datetime
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
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 plot_image(img, figsize=(6,6)):
|
|
fig, ax = plt.subplots(figsize=figsize)
|
|
ax.imshow(img)
|
|
ax.axis("off")
|
|
|
|
return fig
|
|
|
|
|
|
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 |