def main():
if not args.class_map.exists() or not args.class_map.is_file():
eprint("File %s is invalid!" % args.class_map)
with open(args.class_map, 'r') as f:
# Note that we're not using the CLoader since it may not exist on all
# installations, and isn't really necessary for the small file sizes we
# have here.
class_map = yaml.load(f.read(), yaml.Loader)
print("Loaded class mapping:", class_map)
short_classes = set(class_map.keys())
long_classes = set(class_map.values())
if not args.folder.exists() or not args.folder.is_dir():
eprint("Specified folder %s is invalid!" % args.folder)
for root, dirs, files in os.walk(args.folder, followlinks=True):
for name in files:
if not name.endswith(".txt"): continue
if name.endswith("rects.txt"): continue
txt_name = name
txt_full_path = os.path.join(root, txt_name)
bboxes, classes = bbox_writer.read_bboxes(txt_full_path)
for i, c in enumerate(classes):
if c in short_classes:
classes[i] = class_map[c]
if not classes[i] in long_classes:
print("Class name %s in file %s on line %d is invalid!" %
(classes[i], txt_full_path, i + 1))
bbox_writer.write_bboxes(bboxes, classes, txt_full_path)
def save_frame(frame, bboxes, classes, run_path, frame_number):
frame_path = os.path.join(run_path, "%05d.png" % frame_number)
bbox_path = os.path.join(run_path, "%05d.txt" % frame_number)
if not os.path.isfile(frame_path):
print("Saving frame %d to %s" % (frame_number, frame_path))
cv2.imwrite(frame_path, frame)
bbox_writer.write_bboxes(bboxes, classes, bbox_path)
def save_frame(orig, frame, bboxes, classes, run_name, frame_count):
# Scale the bboxes back down by original scale factor.
# This gives us the tight bounding box for the object, rather than the one
# which has been scaled for the tracker.
bboxes = drawing_utils.scale_bboxes(bboxes, 1 / args.scale)
frame_to_draw = orig.copy()
drawing_utils.draw_bboxes(frame_to_draw, bboxes, classes)
show_scaled("Saved Frame", frame_to_draw)
cv2.imwrite(os.path.join(run_name, "%05d.png" % frame_count), orig)
cv2.imwrite(os.path.join(run_name, "rect_%05d.png" % frame_count), frame)
bbox_writer.write_bboxes(bboxes, classes,
os.path.join(run_name, "%05d.txt" % frame_count))
def save_frame(frame, bboxes, classes, run_path, frame_number, validation):
frame_path = os.path.join(run_path, "%05d.png" % frame_number)
bbox_path = os.path.join(run_path, "%05d.txt" % frame_number)
if validation:
if os.path.isfile(frame_path):
bbox_writer.write_bboxes(bboxes, classes, bbox_path)
else:
print("Image %d not found, skipping." % frame_number)
else:
if not os.path.isfile(frame_path):
print("Saving frame %d to %s" % (frame_number, frame_path))
cv2.imwrite(frame_path, frame)
bbox_writer.write_bboxes(bboxes, classes, bbox_path)
def mouse_callback(event, x, y, flags, params):
im = params.copy()
h, w, c = im.shape
x = int(x / WINDOW_SCALE)
y = int(y / WINDOW_SCALE)
if event == cv2.EVENT_LBUTTONDOWN:
points.append(np.array([x, y]))
if len(points) == 1: # Still drawing a rectangle
cv2.rectangle(im, tuple(points[0]), (x, y), (255, 255, 0), 1, 1)
# If the mouse is moved, draw crosshairs
cv2.line(im, (x, 0), (x, h), (255, 0, 0))
cv2.line(im, (0, y), (w, y), (255, 0, 0))
if len(points) == 2: # We've got a rectangle
bbox = np.array([points[0], points[1] - points[0]]).reshape(-1)
bbox = standardize_bbox(bbox)
cls = str(current_class)
bboxes.append(bbox)
classes.append(cls)
points.clear()
# Write the bboxes out to file.
# We get the filename itself, not the full path.
filename = args.filename.name
bbox_filename = bbox_writer.get_bbox_filename(filename)
bbox_path = os.path.join(os.path.dirname(filename), bbox_filename)
bbox_writer.write_bboxes(bboxes, classes, bbox_path)
print("Wrote bboxes to file:", bbox_path)
drawing_utils.draw_bboxes(im, bboxes, classes)
cv2.putText(im, "Current class: %s" % current_class, (100, 20),
cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)
show_scaled(window, im)