# ---------------------------------------------------------------------------
# road_signs.py
#
# AI project on Computer Vision: reading a photo of a road sign and
# identifying its shape - and so what kind of sign it is.
#
# CBSE Class X, Artificial Intelligence (417), Part C project work.
# SDG 3 (Good Health and Well-being - target 3.6, road safety) and
# SDG 11 (Sustainable Cities and Communities).
#
# In India the shape of a road sign tells you its type before you can read it:
#   circle    (red)  - MANDATORY: an order you must obey
#   triangle  (red)  - CAUTIONARY: danger ahead
#   rectangle (blue) - INFORMATORY: useful information
#
# Shape is turned into two numbers (FEATURES):
#   fill  - how much of the box around the sign the sign fills
#           (a rectangle fills almost all of it, a circle about 0.79,
#           a triangle about half)
#   top   - how wide the top of the sign is, compared with its widest row
#           (a triangle is narrow at the top)
# Then two models are compared: a RULE BASED model written from geometry,
# and a LEARNING BASED model that learns the average of each shape.
#
# Libraries: numpy, matplotlib (images and charts), pandas, os.
# ---------------------------------------------------------------------------

# %% Stage 1 - Problem scoping
print("=" * 60)
print("WHAT SHAPE IS THIS SIGN?")
print("=" * 60)
print("Who   : drivers, and the makers of driver-assistance systems")
print("What  : a sign missed at speed can cause an accident")
print("Where : Indian roads, where the shape tells the type of sign")
print("Why   : a car that reads signs can warn a distracted driver")
print()

# %% Stage 2 - Data acquisition
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

pd.set_option("display.width", 120)

shapes = ["circle", "triangle", "rectangle"]
meaning = {"circle": "MANDATORY", "triangle": "CAUTIONARY", "rectangle": "INFORMATORY"}

paths = []
labels = []
for s in shapes:
    for name in sorted(os.listdir(os.path.join("images", s))):
        paths.append(os.path.join("images", s, name))
        labels.append(s)

print("--- Stage 2: Data acquisition ---")
for s in shapes:
    print(s.ljust(10), "photos :", labels.count(s))
print("Total photos      :", len(paths))
print()

# Chart 1: four photos of each shape.
plt.figure(figsize=(7, 5.4))
for r in range(3):
    group = []
    for i in range(len(paths)):
        if labels[i] == shapes[r]:
            group.append(paths[i])
    for c in range(4):
        plt.subplot(3, 4, r * 4 + c + 1)
        plt.imshow(plt.imread(group[c * len(group) // 4]))
        plt.axis("off")
        if c == 0:
            plt.title(shapes[r], fontsize=9, loc="left")
plt.tight_layout()
plt.savefig("chart1_sample_signs.png")
plt.close()

# %% Stage 3 - Data exploration
print("--- Stage 3: Data exploration ---")

# Step 1 - find the sign. Its border is strongly red or strongly blue, which
# the grey and green background never is.
# Step 2 - for every row of the photo, measure the sign's width from its
# leftmost coloured pixel to its rightmost. Adding those widths up gives the
# sign's area, white middle included.
rows = []
for i in range(len(paths)):
    img = plt.imread(paths[i])[:, :, :3].astype(float) * 255
    r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
    red = (r > g + 70) & (r > b + 70)
    blue = (b > r + 60) & (b > g + 40)
    coloured = red | blue
    widths = []
    for y in range(img.shape[0]):
        xs = np.where(coloured[y])[0]
        if len(xs) > 0:
            widths.append(xs[-1] - xs[0] + 1)
    widths = np.array(widths)
    height = len(widths)
    box = height * widths.max()
    fill = widths.sum() / box
    top = widths[: max(1, height // 4)].mean() / widths.max()
    if red.sum() > blue.sum():
        colour = "red"
    else:
        colour = "blue"
    rows.append([paths[i], labels[i], colour, round(fill, 3), round(top, 3)])

data = pd.DataFrame(rows, columns=["photo", "label", "colour", "fill", "top"])
print("Average features of each shape:")
print(data.groupby("label")[["fill", "top"]].mean().round(3).reindex(shapes))
print()
print("Colour found for each shape:")
print(pd.crosstab(data["label"], data["colour"]).reindex(shapes))
print()

# Chart 2: the two features, one dot per photo.
plt.figure(figsize=(7, 5))
dot_colours = ["#d62728", "#ff7f0e", "#1f77b4"]
for k in range(3):
    part = data[data["label"] == shapes[k]]
    plt.scatter(part["fill"], part["top"], color=dot_colours[k], label=shapes[k], alpha=0.75)
plt.title("Two numbers that describe a shape")
plt.xlabel("fill: share of the box the sign fills")
plt.ylabel("top: width of the top / widest width")
plt.legend()
plt.grid(True, linestyle="--", alpha=0.4)
plt.tight_layout()
plt.savefig("chart2_shape_features.png")
plt.close()

# %% Stage 4 - Modelling
print("--- Stage 4: Modelling ---")
use = ["fill", "top"]
data = data.sample(frac=1, random_state=6)
cut = int(len(data) * 0.8)
train = data[:cut]
test = data[cut:]
print("Training photos :", len(train))
print("Testing photos  :", len(test))
print()

# Model A - RULE BASED, from geometry alone, with no training at all:
#   a rectangle fills nearly all of its box; of the rest, a triangle has a
#   narrow top and a circle does not.
rule = []
for i in range(len(test)):
    row = test.iloc[i]
    if row["fill"] > 0.9:
        rule.append("rectangle")
    elif row["top"] < 0.45:
        rule.append("triangle")
    else:
        rule.append("circle")

# Model B - LEARNING BASED: the average fill and top of each shape, learned
# from the training photos; a new sign gets the nearest average.
averages = train.groupby("label")[use].mean()
print("What model B learned - the average of each shape:")
print(averages.round(3).reindex(shapes))
print()
learned = []
for i in range(len(test)):
    distance = ((averages - test[use].iloc[i]) ** 2).sum(axis=1)
    learned.append(distance.idxmin())

# %% Stage 5 - Evaluation
print("--- Stage 5: Evaluation ---")
actual = list(test["label"])
rule_right = 0
learned_right = 0
for i in range(len(actual)):
    if rule[i] == actual[i]:
        rule_right = rule_right + 1
    if learned[i] == actual[i]:
        learned_right = learned_right + 1
print("Rule based     :", rule_right, "of", len(actual), "correct")
print("Learning based :", learned_right, "of", len(actual), "correct")
print("Accuracy of the rule based model     :", round(rule_right / len(actual) * 100, 1), "%")
print("Accuracy of the learning based model :", round(learned_right / len(actual) * 100, 1), "%")
print()

matrix = pd.crosstab(pd.Series(actual, name="Actual"), pd.Series(learned, name="Predicted"))
matrix = matrix.reindex(index=shapes, columns=shapes, fill_value=0)
print("Confusion matrix of the learning based model:")
print(matrix)
print()

wrong = []
for i in range(len(actual)):
    if learned[i] != actual[i] or rule[i] != actual[i]:
        wrong.append(i)
print("Photos at least one model got wrong :", len(wrong))
for i in wrong:
    row = test.iloc[i]
    print("  " + os.path.basename(row["photo"]).ljust(18), "fill", row["fill"], " top", row["top"],
          " rule said", rule[i], " learned said", learned[i])
print()

# Chart 3: confusion matrix.
plt.figure(figsize=(5.5, 4.6))
plt.imshow(matrix.values, cmap="Purples")
for r in range(3):
    for c in range(3):
        value = matrix.values[r][c]
        if value > matrix.values.max() / 2:
            ink = "white"
        else:
            ink = "black"
        plt.text(c, r, value, ha="center", va="center", fontsize=14, color=ink)
plt.xticks(range(3), shapes)
plt.yticks(range(3), shapes)
plt.xlabel("Predicted shape")
plt.ylabel("Actual shape")
plt.title("Learning based model (testing photos)")
plt.tight_layout()
plt.savefig("chart3_matrix.png")
plt.close()

# Chart 4: the photos that fooled a model.
if len(wrong) > 0:
    plt.figure(figsize=(2 * len(wrong) + 0.5, 2.6))
    for n in range(len(wrong)):
        row = test.iloc[wrong[n]]
        plt.subplot(1, len(wrong), n + 1)
        plt.imshow(plt.imread(row["photo"]))
        plt.title("really " + row["label"], fontsize=9)
        plt.axis("off")
    plt.tight_layout()
    plt.savefig("chart4_mistakes.png")
    plt.close()

# %% Stage 6 - Deployment
print("--- Stage 6: Deployment ---")
new_names = sorted(os.listdir("new_photos"))
plt.figure(figsize=(7, 2.8))
for i in range(len(new_names)):
    photo = plt.imread(os.path.join("new_photos", new_names[i]))[:, :, :3]
    img = photo.astype(float) * 255
    r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
    coloured = ((r > g + 70) & (r > b + 70)) | ((b > r + 60) & (b > g + 40))
    widths = []
    for y in range(img.shape[0]):
        xs = np.where(coloured[y])[0]
        if len(xs) > 0:
            widths.append(xs[-1] - xs[0] + 1)
    widths = np.array(widths)
    point = pd.Series({
        "fill": widths.sum() / (len(widths) * widths.max()),
        "top": widths[: max(1, len(widths) // 4)].mean() / widths.max(),
    })
    shape = ((averages - point) ** 2).sum(axis=1).idxmin()
    print(new_names[i], "-> fill", round(point["fill"], 2), " top", round(point["top"], 2),
          "->", shape, "->", meaning[shape], "sign")
    plt.subplot(1, len(new_names), i + 1)
    plt.imshow(photo)
    plt.title(shape + "\n" + meaning[shape].lower(), fontsize=9)
    plt.axis("off")
plt.tight_layout()
plt.savefig("chart5_new_signs.png")
plt.close()