gobot/board-vision/cnn_model/main.py

import cv2
import argparse
import numpy as np
import os
import uuid
import zipfile

from fileSetUtils import compress
from util import process_frame, WINDOW_WIDTH

# Parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("camera", type=int, help="Camera index")
parser.add_argument("output", type=str, help="Output file")
args = parser.parse_args()

cam = cv2.VideoCapture(args.camera)

def convert_degree_to_byte(d):
    return int((d/360)*255)

WINDOW_NAME = "Output"
cv2.namedWindow(WINDOW_NAME)
cv2.moveWindow(WINDOW_NAME, 1,0)

first = True

OUTPUT_FOLDER = args.output

os.system("rm -rf " + OUTPUT_FOLDER)

# Check if folder exists and create it if it doesn't
if not os.path.exists(OUTPUT_FOLDER):
    os.makedirs(OUTPUT_FOLDER)

# Create Postiv and negative folders
if not os.path.exists(OUTPUT_FOLDER + "/positive"):
    os.makedirs(OUTPUT_FOLDER + "/positive")

if not os.path.exists(OUTPUT_FOLDER + "/negative"):
    os.makedirs(OUTPUT_FOLDER + "/negative")

total = 0
neg_total = 0
pos_total = 0
thresh = 0.5

font = cv2.FONT_HERSHEY_SIMPLEX 

try:
    while True:
        ret, frame = cam.read()
        x, y = frame.shape[1], frame.shape[0]

        if first:
            first = False
            print(x, y)   


        # Color Masking Blue Markers

        if not ret:
            print("failed to grab frame")
            break
        
        frame_gray, mask, pos_windows, neg_windows, neutral_windows = process_frame(frame, thresh)
        frame_neutral = cv2.cvtColor(frame_gray, cv2.COLOR_GRAY2BGR)

        for i in neutral_windows:
            cv2.polylines(frame_neutral, [np.array([
                [i[0], i[1]],
                [i[0]+WINDOW_WIDTH, i[1]],
                [i[0]+WINDOW_WIDTH, i[1]+WINDOW_WIDTH],
                [i[0], i[1]+WINDOW_WIDTH]
            ])], True, (255, 128, 0))

        for i in pos_windows:
            cv2.polylines(frame_neutral, [np.array([
                [i[0], i[1]],
                [i[0]+WINDOW_WIDTH, i[1]],
                [i[0]+WINDOW_WIDTH, i[1]+WINDOW_WIDTH],
                [i[0], i[1]+WINDOW_WIDTH]
            ])], True, (0, 255, 255))

            cv2.polylines(frame, [np.array([
                [i[0], i[1]],
                [i[0]+WINDOW_WIDTH, i[1]],
                [i[0]+WINDOW_WIDTH, i[1]+WINDOW_WIDTH],
                [i[0], i[1]+WINDOW_WIDTH]
            ])], True, (0, 255, 255))

        # Tiling 
        res = np.concatenate((
            np.concatenate((
                frame, 
                cv2.cvtColor(frame_gray, cv2.COLOR_GRAY2BGR)            
            ), axis=1),
            np.concatenate((
                cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR), 
                frame_neutral
            ), axis=1),
        ), axis=0)

        res = cv2.putText(res, f"POS: {len(pos_windows)}", (20, 20), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
        res = cv2.putText(res, f"NEG: {len(neg_windows)}", (20, 40), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
        res = cv2.putText(res, f"NEU: {len(neutral_windows)}", (20, 60), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
        res = cv2.putText(res, "THR: {:.2f}".format(thresh), (20, 80), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA)

        cv2.imshow(WINDOW_NAME, res)

        # Wait for ESC key
        k = cv2.waitKey(1)
        if k%256 == 27: # ESC pressed
            print("Escape hit, closing...")
            break

        elif k%256 == 43: # + pressed
            thresh = max(0, thresh + 0.01)

        elif k%256 == 45: # - pressed
            thresh = max(0, thresh - 0.01)

        elif k%256 == 32: # Space pressed
            random_id = str(uuid.uuid4())
            cv2.imwrite(OUTPUT_FOLDER + "/" + random_id + ".png", frame_gray)

            os.mkdir(OUTPUT_FOLDER + f"/positive/{random_id}")
            os.mkdir(OUTPUT_FOLDER + f"/negative/{random_id}")

            # Save negtaive windows 
            for i, window in enumerate(neg_windows):
                cv2.imwrite(OUTPUT_FOLDER + f"/negative/{random_id}/" + str(i) + ".png", frame_gray[window[1]:window[1]+WINDOW_WIDTH, window[0]:window[0]+WINDOW_WIDTH])
            
            # Save positive windows
            for i, window in enumerate(pos_windows):
                cv2.imwrite(OUTPUT_FOLDER + f"/positive/{random_id}/" + str(i) + ".png", frame_gray[window[1]:window[1]+WINDOW_WIDTH, window[0]:window[0]+WINDOW_WIDTH])
            
            count_neg = len(neg_windows)
            count_pos = len(pos_windows)

            total += count_neg + count_pos
            neg_total += count_neg
            pos_total += count_pos

            print("+{:03d} -{:07d} ?{:04d} ({:07d}) | +{:07d} -{:07d} ({:09d})".format(
                count_pos, count_neg, len(neutral_windows), len(neutral_windows) + count_neg + count_pos, 
                pos_total, neg_total, total))
        

except KeyboardInterrupt:
    print("closing...")

# Check if output_image archive exists
if os.path.exists(args.output + "_images.zip"):
    os.remove(args.output + "_images.zip")

# Create archive of images
image_archive = zipfile.ZipFile(args.output + "_images.zip", 'w')

# Iterate over all images in the output folder
for files in os.listdir(OUTPUT_FOLDER):
    if files.endswith(".png"):
        image_archive.write(OUTPUT_FOLDER + "/" + files, files)
    
# Delete output zip file if it exists
if os.path.exists(args.output + ".zip"):
    os.remove(args.output + ".zip")

# Compress the output folder
compress(args.output + ".zip", OUTPUT_FOLDER)

os.system("rm -rf " + OUTPUT_FOLDER)

cam.release()
cv2.destroyAllWindows()