nicknochnack / TFODCourse

944 stars 919 forks source link

Image capture from camera on button click #136

Open Sinet2000 opened 1 year ago

Sinet2000 commented 1 year ago

There is an upgraded code, that could be useful to others. I am using both USB and an integrated laptop camera. In the upgraded code, the image capture screen with the label will appear, where the user must click t to capture an image


def find_camera_index():
    index = 0
    while True:
        cap = cv2.VideoCapture(index, cv2.CAP_DSHOW)
        if not cap.read()[0]:
            break
        cap.release()
        index += 1
    return index - 1

def take_photo(cap, label):
    ret, frame = cap.read()
    imgname = os.path.join(IMAGES_PATH, label, label + '.' + '{}.jpg'.format(str(uuid.uuid1())))
    cv2.imwrite(imgname, frame)
    return frame

def draw_label(frame, label):
    font = cv2.FONT_HERSHEY_SIMPLEX
    position = (50, 50)
    font_scale = 1
    color = (255, 0, 0)
    thickness = 2
    return cv2.putText(frame, label, position, font, font_scale, color, thickness)

camera_index = find_camera_index()
cap = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW)

for label in labels:
    print('Collecting images for {}'.format(label))
    time.sleep(1)

    for imgnum in range(number_imgs):
        print('Press "t" to take a photo for {}'.format(label))
        while True:
            ret, frame = cap.read()
            labeled_frame = draw_label(frame, label)
            cv2.imshow('frame', labeled_frame)
            key = cv2.waitKey(1) & 0xFF

            if key == ord('t'):
                taken_frame = take_photo(cap, label)
                print('Photo {} for {} has been taken'.format(imgnum, label))
                break
            elif key == ord('q'):
                break

        if key == ord('q'):
            break

cap.release()
cv2.destroyAllWindows()