175 lines
5.4 KiB (Stored with Git LFS)
Python
175 lines
5.4 KiB (Stored with Git LFS)
Python
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() |