pre_labeling.py

"""
Copyright (c) 2025 Samsung Electronics Co., Ltd.

Author(s):
Mahmoud Afifi (m.afifi1@samsung.com, m.3afifi@gmail.com)

Licensed under the Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) License, (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at https://creativecommons.org/licenses/by-nc/4.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and limitations under the License.
For conditions of distribution and use, see the accompanying LICENSE.md file.

This script is used to extract metadata and raw images from DNG files. The extracted files are then processed using a
Matlab GUI for annotation to get the final dataset files.
"""

import utils
import argparse
import os
import shutil
import numpy as np

def get_args():
  parser = argparse.ArgumentParser(
    description='Process and extract data from DNG and JPG images into structured output folders.')
  parser.add_argument('-dp', '--dataset_path', type=str, required=True,
                      help='Path to the directory containing the DNG and JPG images extracted from the camera.')
  parser.add_argument('-op', '--output_path', type=str, required=True,
                      help='Path to the output directory where processed files will be saved.')
  parser.add_argument('-rd', '--raw_dir', type=str, default='raw_images',
                      help='Name of the subfolder within the output directory to save processed raw images.')
  parser.add_argument('-md', '--metadata_dir', type=str, default='data',
                      help='Name of the subfolder within the output directory to save extracted metadata files.')
  parser.add_argument('-dd', '--dng_dir', type=str, default='dngs',
                      help='Name of the subfolder within the dataset path that contains the DNG files.')
  parser.add_argument('-sd', '--srgb_dir', type=str, default='srgb_images',
                      help='Name of the subfolder within the dataset path that contains the sRGB images.')
  return parser.parse_args()

if __name__ == '__main__':
  args = get_args()
  dng_files = [f for f in os.listdir(args.dataset_path) if f.endswith('.dng') or f.endswith('.DNG')]
  pre_lat_lon = None
  os.makedirs(args.output_path, exist_ok=True)
  sub_dirs = [args.raw_dir, args.metadata_dir, args.dng_dir, args.srgb_dir]
  for sub_dir in sub_dirs:
    os.makedirs(os.path.join(args.output_path, sub_dir), exist_ok=True)

  for i, dng_file in enumerate(dng_files):
    print(f'Processing {i+1}/{len(dng_files)}...')
    try:
      dng_path = os.path.join(args.dataset_path, dng_file)
      filename = os.path.splitext(dng_file)[0]
      raw = utils.extract_image_from_dng(dng_path)
      dng_metadata = utils.extract_raw_metadata(dng_path)
      normalized_raw = utils.normalize_raw(raw, black_level=dng_metadata['black_level'],
                                           white_level=dng_metadata['white_level'])
      normalized_raw = utils.demosaice(normalized_raw, cfa_pattern=dng_metadata['pattern'])
      jpg_metadata, pre_lat_lon, _ = utils.extract_jpg_metadata(dng_path.replace('.dng', '.jpg'),
                                                                pre_lat_lon=pre_lat_lon)
      gw_illum = np.mean(np.reshape(normalized_raw, [-1, 3]), axis=0)
      utils.imwrite(normalized_raw, os.path.join(str(args.output_path), args.raw_dir, filename), 'PNG-16')
      metadata = {'cam_illum': dng_metadata['illum_color'], 'cam_daylight_illum': dng_metadata['daylight_illum_color'],
                  'ccm': dng_metadata['color_matrix'], 'capture_metadata': jpg_metadata, 'gw_illum': gw_illum.tolist()}
      utils.write_json_file(metadata, os.path.join(str(args.output_path), args.metadata_dir, filename))
      shutil.move(dng_path, os.path.join(str(args.output_path), args.dng_dir, dng_file))
      shutil.move(dng_path.replace('.dng', '.jpg'),
                  os.path.join(str(args.output_path), args.srgb_dir, dng_file.replace('.dng', '.jpg')))
    except ValueError:
      print(f"Couldn't process image # {i+1}")