yoco_first_version / preprocess.py
nathbns's picture
Update preprocess.py
b6bd9ea verified
import numpy as np
import cv2
import glob
from pathlib import Path
from matplotlib import pyplot as plt
from rescale import *
from slid import detect_lines
from laps import LAPS
from llr import LLR, llr_pad
RAW_DATA_FOLDER = './data/raw/games/'
PREPROCESSED_FOLDER = './data/preprocessed/games/'
def preprocess_image(path, final_folder="", filename="", save=False):
res = cv2.imread(path)[..., ::-1]
for _ in range(2):
img, shape, scale = image_resize(res)
lines = detect_lines(img)
lattice_points = LAPS(img, lines)
inner_points = LLR(img, lattice_points, lines)
four_points = llr_pad(inner_points, img)
try:
res = crop(res, four_points, scale)
except:
print("WARNING: couldn't crop around outer points")
res = crop(
res, inner_points, scale)
if save:
Path(final_folder).mkdir(parents=True, exist_ok=True)
plt.imsave("%s/%s" % (final_folder, filename), res)
return res
def preprocess_single_image(game_name, version, image_number):
raw_path = RAW_DATA_FOLDER + '%s/%s/%i.*' % (game_name, version, image_number)
final_folder = PREPROCESSED_FOLDER + "%s/%s/" % (game_name, version)
matching_files = glob.glob(raw_path)
if not matching_files:
print(f"ERROR: No image found at {raw_path}")
return False
if len(matching_files) > 1:
print(f"WARNING: Multiple files found, using first one: {matching_files[0]}")
path = matching_files[0]
print(f"Processing: {path}")
try:
preprocess_image(path, final_folder=final_folder,
filename="%i.png" % image_number, save=True)
print(f"βœ“ Successfully saved to {final_folder}{image_number}.png")
return True
except Exception as e:
print(f"βœ— ERROR processing image: {e}")
import traceback
traceback.print_exc()
return False
def preprocess_multiple_images(images_list):
total = len(images_list)
success = 0
failed = []
print(f"Processing {total} images...")
for game_name, version, img_num in images_list:
print(f"\n[{success + len(failed) + 1}/{total}] {game_name}/{version}/image {img_num}")
if preprocess_single_image(game_name, version, img_num):
success += 1
else:
failed.append((game_name, version, img_num))
print(f"\n{'='*50}")
print(f"Summary: {success}/{total} images processed successfully")
if failed:
print(f"Failed images ({len(failed)}):")
for game, ver, num in failed:
print(f" - {game}/{ver}/image {num}")
def preprocess_games(game_list):
for game_name in game_list:
for ver in ['orig', 'rev']:
img_filename_list = []
folder_name = RAW_DATA_FOLDER + '%s/%s/*' % (game_name, ver)
for path_name in glob.glob(folder_name):
img_filename_list.append(path_name)
final_folder = PREPROCESSED_FOLDER + "%s/%s/" % (game_name, ver)
count = 0
img_filename_list.sort(key=lambda s: int(
s.split('/')[-1].split('.')[0]))
for path in img_filename_list:
count += 1
preprocess_image(path, final_folder=final_folder,
filename="%i.png" % count, save=True)
if count > 0:
print("Done saving in %s." % final_folder)
else:
print("No images found in %s" % folder_name)
if __name__ == '__main__':
# deja fait: carlsen_anand_2014, carlsen_gukesh_2025 vachier-lagrave_carlsen_2023 david_vachier-lagrave_2014
# game_list = ['david_vachier-lagrave_2014']
# preprocess_games(game_list)
preprocess_single_image('vachier-lagrave_carlsen_2023', 'rev', 57)