diff --git a/.github/workflows/python_pycsou.yml b/.github/workflows/python_pycsou.yml index 61f89fa5..dfdafb9b 100644 --- a/.github/workflows/python_pycsou.yml +++ b/.github/workflows/python_pycsou.yml @@ -20,7 +20,7 @@ jobs: fail-fast: false max-parallel: 12 matrix: - os: [ubuntu-latest, macos-latest, windows-latest] + os: [ubuntu-latest, macos-12, windows-latest] python-version: [3.9, "3.10"] steps: - uses: actions/checkout@v3 diff --git a/configs/collect_dataset.yaml b/configs/collect_dataset.yaml index c54ef8e7..75df8597 100644 --- a/configs/collect_dataset.yaml +++ b/configs/collect_dataset.yaml @@ -1,6 +1,6 @@ # python scripts/collect_dataset_on_device.py -cn collect_dataset -input_dir: /mnt/mirflickr/10 +input_dir: /mnt/mirflickr/all input_file_ext: jpg # can pass existing folder to continue measurement @@ -41,6 +41,9 @@ display: landscape: False # whether to force landscape capture: + measure_bg: False # measure bg every x images, set False if not measuring background + bg_fp: "black_background" + framerate: 10 skip: False # to test looping over displaying images config_pause: 3 iso: 100 diff --git a/configs/collect_dataset_background.yaml b/configs/collect_dataset_background.yaml new file mode 100644 index 00000000..9f509eda --- /dev/null +++ b/configs/collect_dataset_background.yaml @@ -0,0 +1,26 @@ +# python scripts/measure/collect_dataset_on_device.py -cn collect_dataset_background +defaults: + - collect_dataset + - _self_ + + +output_dir: /mnt/mirflickr/all_measured_20240813-183259 + +# files to measure +n_files: 25000 + +min_level: 160 +max_tries: 3 + + +# -- display parameters +display: + screen_res: [1920, 1200] # width, height + image_res: [600, 600] # useful if input images don't have the same dimension, set it to this + vshift: -34 + +capture: + measure_bg: 1 # measure bg every x images, set False if not measuring background + awb_gains: [1.8, 1.1] # red, blue + fact_increase: 1.35 # multiplicative factor to increase exposure + fact_decrease: 1.3 diff --git a/scripts/measure/collect_dataset_on_device.py b/scripts/measure/collect_dataset_on_device.py index 8bb34077..4b1314b2 100644 --- a/scripts/measure/collect_dataset_on_device.py +++ b/scripts/measure/collect_dataset_on_device.py @@ -11,6 +11,7 @@ To test on local machine, set dummy=True (which will just copy the files over). """ +import json import numpy as np import hydra @@ -28,7 +29,6 @@ import glob from lensless.hardware.slm import set_programmable_mask, adafruit_sub2full - from lensless.hardware.constants import ( RPI_HQ_CAMERA_CCM_MATRIX, RPI_HQ_CAMERA_BLACK_LEVEL, @@ -52,7 +52,6 @@ def natural_sort(arr): @hydra.main(version_base=None, config_path="../../configs", config_name="collect_dataset") def collect_dataset(config): - input_dir = config.input_dir output_dir = config.output_dir if output_dir is None: @@ -162,7 +161,7 @@ def collect_dataset(config): camera.close() # -- now set up camera with desired settings - camera = PiCamera(sensor_mode=0, resolution=tuple(res)) + camera = PiCamera(sensor_mode=0, resolution=tuple(res), framerate=config.capture.framerate) # Set ISO to the desired value camera.resolution = tuple(res) @@ -202,6 +201,10 @@ def collect_dataset(config): exposure_vals = [] brightness_vals = [] n_tries_vals = [] + + bg_name = None + current_bg = {} + shutter_speed = init_shutter_speed for i, _file in enumerate(tqdm.tqdm(files[start_idx:])): # save file in output directory as PNG @@ -210,147 +213,55 @@ def collect_dataset(config): # if not done, perform measurement if not os.path.isfile(output_fp): - if config.dummy: shutil.copyfile(_file, output_fp) time.sleep(1) - else: - # -- show on display - screen_res = np.array(config.display.screen_res) - hshift = config.display.hshift - vshift = config.display.vshift - pad = config.display.pad - brightness = init_brightness - display_image_path = config.display.output_fp - rot90 = config.display.rot90 - display_command = f"python scripts/measure/prep_display_image.py --fp {_file} --output_path {display_image_path} --screen_res {screen_res[0]} {screen_res[1]} --hshift {hshift} --vshift {vshift} --pad {pad} --brightness {brightness} --rot90 {rot90}" - if config.display.landscape: - display_command += " --landscape" - if config.display.image_res is not None: - display_command += ( - f" --image_res {config.display.image_res[0]} {config.display.image_res[1]}" - ) - # print(display_command) - os.system(display_command) - - time.sleep(config.display.delay) - - if not config.capture.skip: - - # -- set mask pattern - if config.masks is not None: - mask_idx = (i + start_idx) % config.masks.n - mask_fp = mask_dir / f"mask_{mask_idx}.npy" - print("using mask: ", mask_fp) - mask_vals = np.load(mask_fp) - full_pattern = adafruit_sub2full( - mask_vals, - center=config.masks.center, - ) - set_programmable_mask(full_pattern, device=config.masks.device) - - # -- take picture - max_pixel_val = 0 - fact_increase = config.capture.fact_increase - fact_decrease = config.capture.fact_decrease - n_tries = 0 - - camera.shutter_speed = init_shutter_speed - time.sleep(config.capture.config_pause) + # display img + display_img(_file, config, init_brightness) + # capture img + output, _, _, camera = capture_screen(MAX_LEVEL, MAX_TRIES, MIN_LEVEL, _file, brightness_vals, camera, config, down, + exposure_vals, g, i, init_brightness, shutter_speed, None, + n_tries_vals, + output_fp, start_idx) + + if config.capture.measure_bg: + # name of background for current image + bg_name = plib.Path(config.capture.bg_fp + str(i)).with_suffix(f".{config.output_file_ext}") + bg = output_dir / bg_name + + # append current file to bg list + if str(bg_name) not in current_bg: + current_bg[str(bg_name)] = str(_file.name) + else: + current_bg[str(bg_name)].append(str(_file.name)) + # capture background periodically + if i % config.capture.measure_bg == 0 or (i == n_files - 1): + bg_name = str(bg_name) + # push the last bg-capture pairs + if current_bg: + with open(output_dir / "bg_mappings.json", 'a') as outfile: + json.dump(current_bg, outfile, indent=4) + current_bg = {} + #current_bg[bg_name] = None + + # display bg + display_img(None, config, brightness=init_brightness) + # capture bg + output, shutter_speed, init_brightness, camera = capture_screen(MAX_LEVEL, 0, MIN_LEVEL, + plib.Path(config.capture.bg_fp + str(i)).with_suffix(f".{config.output_file_ext}"), + brightness_vals, camera, config, down, + exposure_vals, g, i, init_brightness, shutter_speed, None, + n_tries_vals, + bg, start_idx) - current_screen_brightness = init_brightness - current_shutter_speed = camera.shutter_speed - print(f"current shutter speed: {current_shutter_speed}") - print(f"current screen brightness: {current_screen_brightness}") - - while max_pixel_val < MIN_LEVEL or max_pixel_val > MAX_LEVEL: - - # get bayer data - stream = picamerax.array.PiBayerArray(camera) - camera.capture(stream, "jpeg", bayer=True) - output_bayer = np.sum(stream.array, axis=2).astype(np.uint16) - - # convert to RGB - output = bayer2rgb_cc( - output_bayer, - down=down, - nbits=12, - blue_gain=float(g[1]), - red_gain=float(g[0]), - black_level=RPI_HQ_CAMERA_BLACK_LEVEL, - ccm=RPI_HQ_CAMERA_CCM_MATRIX, - nbits_out=8, - ) - - # if down: - # output = resize( - # output[None, ...], factor=1 / down, interpolation=cv2.INTER_CUBIC - # )[0] - - # save image - save_image(output, output_fp, normalize=False) - - # print range - print(f"{output_fp}, range: {output.min()} - {output.max()}") - - n_tries += 1 - if n_tries > MAX_TRIES: - print("Max number of tries reached!") - break - - max_pixel_val = output.max() - if max_pixel_val < MIN_LEVEL: - - # increase exposure - current_shutter_speed = int(current_shutter_speed * fact_increase) - camera.shutter_speed = current_shutter_speed - time.sleep(config.capture.config_pause) - - print(f"increasing shutter speed to [desired] {current_shutter_speed} [actual] {camera.shutter_speed}") - - elif max_pixel_val > MAX_LEVEL: - - if current_shutter_speed > 13098: # TODO: minimum for RPi HQ - # decrease exposure - current_shutter_speed = int(current_shutter_speed / fact_decrease) - camera.shutter_speed = current_shutter_speed - time.sleep(config.capture.config_pause) - print(f"decreasing shutter speed to [desired] {current_shutter_speed} [actual] {camera.shutter_speed}") - - else: - - # decrease screen brightness - current_screen_brightness = current_screen_brightness - 10 - screen_res = np.array(config.display.screen_res) - hshift = config.display.hshift - vshift = config.display.vshift - pad = config.display.pad - brightness = current_screen_brightness - display_image_path = config.display.output_fp - rot90 = config.display.rot90 - - display_command = f"python scripts/measure/prep_display_image.py --fp {_file} --output_path {display_image_path} --screen_res {screen_res[0]} {screen_res[1]} --hshift {hshift} --vshift {vshift} --pad {pad} --brightness {brightness} --rot90 {rot90}" - if config.display.landscape: - display_command += " --landscape" - if config.display.image_res is not None: - display_command += f" --image_res {config.display.image_res[0]} {config.display.image_res[1]}" - # print(display_command) - os.system(display_command) - - time.sleep(config.display.delay) - - exposure_vals.append(current_shutter_speed / 1e6) - brightness_vals.append(current_screen_brightness) - n_tries_vals.append(n_tries) if recon is not None: - # normalize and remove background output = output.astype(np.float32) output /= output.max() - output -= bg + output -= bg # TODO implement fancy bg subtraction output = np.clip(output, a_min=0, a_max=output.max()) # set data @@ -366,11 +277,11 @@ def collect_dataset(config): if config.runtime: proc_time = time.time() - start_time if proc_time > runtime_sec: - print(f"-- measured {i+1} / {n_files} files") + print(f"-- measured {i + 1} / {n_files} files") break proc_time = time.time() - start_time - print(f"\nFinished, {proc_time/60.:.3f} minutes.") + print(f"\nFinished, {proc_time / 60.:.3f} minutes.") # print brightness and exposure range and average print(f"brightness range: {np.min(brightness_vals)} - {np.max(brightness_vals)}") @@ -381,5 +292,130 @@ def collect_dataset(config): print(f"n_tries average: {np.mean(n_tries_vals)}") +def capture_screen(MAX_LEVEL, MAX_TRIES, MIN_LEVEL, _file, brightness_vals, camera, config, down, exposure_vals, g, i, + init_brightness, init_shutter_speed, mask_dir, n_tries_vals, output_fp, start_idx): + if not config.capture.skip: + + # -- set mask pattern + if config.masks is not None: + mask_idx = (i + start_idx) % config.masks.n + mask_fp = mask_dir / f"mask_{mask_idx}.npy" + print("using mask: ", mask_fp) + mask_vals = np.load(mask_fp) + full_pattern = adafruit_sub2full( + mask_vals, + center=config.masks.center, + ) + set_programmable_mask(full_pattern, device=config.masks.device) + + # -- take picture + max_pixel_val = 0 + fact_increase = config.capture.fact_increase + fact_decrease = config.capture.fact_decrease + n_tries = 0 + + camera.shutter_speed = int(init_shutter_speed) + time.sleep(config.capture.config_pause) + current_shutter_speed = camera.shutter_speed + + current_screen_brightness = init_brightness + print(f"current shutter speed: {current_shutter_speed}") + print(f"current screen brightness: {current_screen_brightness}") + + while max_pixel_val < MIN_LEVEL or max_pixel_val > MAX_LEVEL: + + # get bayer data + stream = picamerax.array.PiBayerArray(camera) + camera.capture(stream, "jpeg", bayer=True) + output_bayer = np.sum(stream.array, axis=2).astype(np.uint16) + + # convert to RGB + output = bayer2rgb_cc( + output_bayer, + down=down, + nbits=12, + blue_gain=float(g[1]), + red_gain=float(g[0]), + black_level=RPI_HQ_CAMERA_BLACK_LEVEL, + ccm=RPI_HQ_CAMERA_CCM_MATRIX, + nbits_out=8, + ) + + # if down: + # output = resize( + # output[None, ...], factor=1 / down, interpolation=cv2.INTER_CUBIC + # )[0] + + # save image + save_image(output, output_fp, normalize=False) + + # print range + print(f"{output_fp}, range: {output.min()} - {output.max()}") + + n_tries += 1 + if n_tries > MAX_TRIES: + if MAX_TRIES != 0: + print("Max number of tries reached!") + break + + max_pixel_val = output.max() + + if max_pixel_val < MIN_LEVEL: + + # increase exposure + current_shutter_speed = int(current_shutter_speed * fact_increase) + camera.shutter_speed = current_shutter_speed + time.sleep(config.capture.config_pause) + + print(f"increasing shutter speed to [desired] {current_shutter_speed} [actual] {camera.shutter_speed}") + + elif max_pixel_val > MAX_LEVEL: + if current_shutter_speed > 13098: # TODO: minimum for RPi HQ + # decrease exposure + current_shutter_speed = int(current_shutter_speed / fact_decrease) + camera.shutter_speed = current_shutter_speed + time.sleep(config.capture.config_pause) + print(f"decreasing shutter speed to [desired] {current_shutter_speed} [actual] {camera.shutter_speed}") + + else: + + # decrease screen brightness + current_screen_brightness = current_screen_brightness - 10 + display_img(_file, config, current_screen_brightness) + + exposure_vals.append(current_shutter_speed / 1e6) + brightness_vals.append(current_screen_brightness) + n_tries_vals.append(n_tries) + return output, current_shutter_speed, current_screen_brightness, camera + + +def display_img(_file, config, brightness): + if _file is None: + point_source = np.zeros(tuple(config.display.screen_res) + (3,)) + fp = "tmp_display.png" + im = Image.fromarray(point_source.astype("uint8"), "RGB") + im.save(fp) + _file = fp + + # -- show on display + screen_res = np.array(config.display.screen_res) + hshift = config.display.hshift + vshift = config.display.vshift + pad = config.display.pad + + display_image_path = config.display.output_fp + rot90 = config.display.rot90 + display_command = f"python scripts/measure/prep_display_image.py --fp {_file} --output_path {display_image_path} --screen_res {screen_res[0]} {screen_res[1]} --hshift {hshift} --vshift {vshift} --pad {pad} --brightness {brightness} --rot90 {rot90}" + if config.display.landscape: + display_command += " --landscape" + if config.display.image_res is not None: + display_command += ( + f" --image_res {config.display.image_res[0]} {config.display.image_res[1]}" + ) + # print(display_command) + os.system(display_command) + time.sleep(config.display.delay) + + if __name__ == "__main__": collect_dataset()