1#!/usr/bin/env python3 2# -*- coding: utf-8 -*- 3# 4# Copyright 2012-2016 Torsten Bronger <bronger@physik.rwth-aachen.de> 5 6missing_packages = set() 7 8import subprocess, os, os.path, sys, multiprocessing, math, re, contextlib, glob, codecs, struct, argparse 9try: 10 import numpy 11except ImportError: 12 missing_packages.add("python3-numpy") 13try: 14 from scipy.optimize.minpack import leastsq 15except ImportError: 16 missing_packages.add("python3-scipy") 17def test_program(program, package_name): 18 try: 19 subprocess.call([program], stdout=open(os.devnull, "w"), stderr=open(os.devnull, "w")) 20 except OSError: 21 missing_packages.add(package_name) 22test_program("dcraw", "dcraw") 23test_program("convert", "imagemagick") 24test_program("tca_correct", "hugin-tools") 25test_program("exiv2", "exiv2") 26if missing_packages: 27 print("The following packages are missing (Ubuntu packages, names may differ on other systems):\n {0}\nAbort.". 28 format(" ".join(missing_packages))) 29 sys.exit() 30try: 31 dcraw_version = float(subprocess.Popen(["dcraw"], stdout=subprocess.PIPE).communicate()[0].splitlines()[1]. 32 rpartition("v")[2]) 33except: 34 dcraw_version = 0 35h_option = [] if 8.99 < dcraw_version < 9.18 else ["-h"] 36 37 38parser = argparse.ArgumentParser(description="Calibration generator for Lensfun.") 39parser.add_argument("--complex-tca", action="store_true", help="switches on non-linear polynomials for TCA") 40args = parser.parse_args() 41 42 43@contextlib.contextmanager 44def chdir(dirname=None): 45 curdir = os.getcwd() 46 try: 47 if dirname is not None: 48 os.chdir(dirname) 49 yield 50 finally: 51 os.chdir(curdir) 52 53 54class Lens(object): 55 56 def __init__(self, name, maker, mount, cropfactor, aspect_ratio, type_): 57 self.name, self.maker, self.mount, self.cropfactor, self.aspect_ratio, self.type_ = \ 58 name, maker, mount, cropfactor, aspect_ratio, type_ 59 self.calibration_lines = [] 60 self.minimal_focal_length = float("inf") 61 62 def add_focal_length(self, focal_length): 63 self.minimal_focal_length = min(self.minimal_focal_length, focal_length) 64 65 def __lt__(self, other): 66 return self.minimal_focal_length < other.minimal_focal_length 67 68 def write(self, outfile): 69 type_line = " <type>{0}</type>\n".format(self.type_) if self.type_ else "" 70 outfile.write(""" 71 <lens> 72 <maker>{0}</maker> 73 <model>{1}</model> 74 <mount>{2}</mount> 75 <cropfactor>{3}</cropfactor> 76""".format(self.maker, self.name, self.mount, self.cropfactor)) 77 if self.type_: 78 outfile.write(" <type>{0}</type>\n".format(self.type_)) 79 if self.aspect_ratio and self.aspect_ratio != "3:2": 80 outfile.write(" <aspect-ratio>{0}</aspect-ratio>\n".format(self.aspect_ratio)) 81 if self.calibration_lines: 82 outfile.write(" <calibration>\n") 83 for line in self.calibration_lines: 84 outfile.write(" {0}\n".format(line)) 85 outfile.write(" </calibration>\n") 86 outfile.write(" </lens>\n") 87 88 89def generate_raw_conversion_call(filename, dcraw_options): 90 basename, extension = os.path.splitext(filename) 91 extension = extension[1:] 92 if extension.lower() in ["jpg", "jpeg", "tif"]: 93 result = ["convert", filename] 94 if "-4" in dcraw_options: 95 result.extend(["-colorspace", "RGB", "-depth", "16"]) 96 result.append("tiff:-" if "-c" in dcraw_options else basename + ".tiff") 97 return result 98 else: 99 return ["dcraw", "-T", "-t", "0"] + dcraw_options + [filename] 100 101 102raw_file_extensions = ["3fr", "ari", "arw", "bay", "crw", "cr2", "cap", "dcs", "dcr", "dng", "drf", "eip", "erf", "fff", 103 "iiq", "k25", "kdc", "mef", "mos", "mrw", "nef", "nrw", "obm", "orf", "pef", "ptx", "pxn", "r3d", 104 "raf", "raw", "rwl", "rw2", "rwz", "sr2", "srf", "srw", "x3f", "jpg", "jpeg", "tif"] 105def find_raw_files(): 106 result = [] 107 for file_extension in raw_file_extensions: 108 result.extend(glob.glob("*." + file_extension)) 109 result.extend(glob.glob("*." + file_extension.upper())) 110 return result 111 112 113# 114# Collect EXIF data 115# 116 117file_exif_data = {} 118filepath_pattern = re.compile(r"(?P<lens_model>.+)--(?P<focal_length>[0-9.]+)mm--(?P<aperture>[0-9.]+)") 119missing_lens_model_warning_printed = False 120exiv2_candidates = [] 121 122def browse_directory(directory): 123 if os.path.exists(directory): 124 with chdir(directory): 125 for filename in find_raw_files(): 126 full_filename = os.path.join(directory, filename) 127 match = filepath_pattern.match(os.path.splitext(filename)[0]) 128 if match: 129 file_exif_data[full_filename] = \ 130 (match.group("lens_model").replace("##", "=").replace("++", "*").replace("___", ":").replace("__", "/"). 131 replace("_", " "), float(match.group("focal_length")), float(match.group("aperture"))) 132 else: 133 exiv2_candidates.append(full_filename) 134browse_directory("distortion") 135browse_directory("tca") 136for directory in glob.glob("vignetting*"): 137 browse_directory(directory) 138candidates_per_group = len(exiv2_candidates) // multiprocessing.cpu_count() + 1 139candidate_groups = [] 140while exiv2_candidates: 141 candidate_group = exiv2_candidates[:candidates_per_group] 142 if candidate_group: 143 candidate_groups.append(candidate_group) 144 del exiv2_candidates[:candidates_per_group] 145def call_exiv2(candidate_group): 146 exiv2_process = subprocess.Popen( 147 ["exiv2", "-PEkt", "-g", "Exif.Photo.LensModel", "-g", "Exif.Photo.FocalLength", "-g", "Exif.Photo.FNumber"] 148 + candidate_group, stdout=subprocess.PIPE) 149 lines = exiv2_process.communicate()[0].decode("utf-8").splitlines() 150 assert exiv2_process.returncode in [0, 253] 151 result = {} 152 for line in lines: 153 filename, data = line.split("Exif.Photo.") 154 filename = filename.rstrip() 155 if not filename: 156 assert len(candidate_group) == 1 157 filename = candidate_group[0] 158 fieldname, field_value = data.split(None, 1) 159 exif_data = result.setdefault(filename, [None, float("nan"), float("nan")]) 160 if fieldname == "LensModel": 161 exif_data[0] = field_value 162 elif fieldname == "FocalLength": 163 exif_data[1] = float(field_value.partition("mm")[0]) 164 elif fieldname == "FNumber": 165 exif_data[2] = float(field_value.partition("F")[2]) 166 for filename, exif_data in result.copy().items(): 167 result[filename] = tuple(exif_data) 168 return result 169pool = multiprocessing.Pool() 170for group_exif_data in pool.map(call_exiv2, candidate_groups): 171 file_exif_data.update(group_exif_data) 172pool.close() 173pool.join() 174for filename in list(file_exif_data): # list() because I change the dict during iteration 175 lens_model, focal_length, aperture = file_exif_data[filename] 176 if not lens_model: 177 lens_model = "Standard" 178 if not missing_lens_model_warning_printed: 179 print(filename + ":") 180 print("""I couldn't detect the lens model name and assumed "Standard". 181For cameras without interchangable lenses, this may be correct. 182However, this fails if there is data of different undetectable lenses. 183A newer version of exiv2 (if available) may help. 184(This message is printed only once.)\n""") 185 missing_lens_model_warning_printed = True 186 file_exif_data[filename] = (lens_model, focal_length, aperture) 187 if numpy.isnan(focal_length) or numpy.isnan(aperture): 188 print(filename + ":") 189 print("""Aperture and/or focal length EXIF data is missing in this RAW file. 190You have to rename them according to the scheme "Lens_name--16mm--1.4.RAW" 191(Use your RAW file extension of course.) Abort.""") 192 sys.exit() 193 194 195# 196# Generation TIFFs from distortion RAWs 197# 198 199if os.path.exists("distortion"): 200 with chdir("distortion"): 201 pool = multiprocessing.Pool() 202 results = set() 203 for filename in find_raw_files(): 204 if not os.path.exists(os.path.splitext(filename)[0] + ".tiff"): 205 results.add(pool.apply_async(subprocess.call, [generate_raw_conversion_call(filename, ["-w"])])) 206 pool.close() 207 pool.join() 208 [result.get() for result in results] 209 210 211# 212# Parse/generate lenses.txt 213# 214 215lens_line_pattern = re.compile( 216 r"(?P<name>.+):\s*(?P<maker>[^,]+)\s*,\s*(?P<mount>[^,]+)\s*,\s*(?P<cropfactor>[^,]+)" 217 r"(\s*,\s*(?P<aspect_ratio>\d+:\d+|[0-9.]+))?(\s*,\s*(?P<type>[^,]+))?") 218distortion_line_pattern = re.compile(r"\s*distortion\((?P<focal_length>[.0-9]+)mm\)\s*=\s*" 219 r"(?P<a>[-.0-9e]+)(?:\s*,\s*(?P<b>[-.0-9e]+)\s*,\s*(?P<c>[-.0-9e]+))?") 220lenses = {} 221try: 222 for linenumber, original_line in enumerate(open("lenses.txt")): 223 linenumber += 1 224 line = original_line.strip() 225 if line and not line.startswith("#"): 226 match = lens_line_pattern.match(line) 227 if match: 228 data = match.groupdict() 229 current_lens = Lens(data["name"], data["maker"], data["mount"], data["cropfactor"], 230 data["aspect_ratio"], data["type"]) 231 lenses[data["name"]] = current_lens 232 else: 233 match = distortion_line_pattern.match(line) 234 if not match: 235 print("Invalid line {0} in lenses.txt:\n{1}Abort.".format(linenumber, original_line)) 236 sys.exit() 237 data = list(match.groups()) 238 data[0] = float(data[0]) 239 current_lens.add_focal_length(data[0]) 240 if data[2] is None: 241 current_lens.calibration_lines.append( 242 """<distortion model="poly3" focal="{0:g}" k1="{1}"/>""".format(data[0], data[1])) 243 else: 244 current_lens.calibration_lines.append( 245 """<distortion model="ptlens" focal="{0:g}" a="{1}" b="{2}" c="{3}"/>""".format(*data)) 246except IOError: 247 focal_lengths = {} 248 for exif_data in file_exif_data.values(): 249 focal_lengths.setdefault(exif_data[0], set()).add(exif_data[1]) 250 lens_names_by_focal_length = sorted((min(lengths), lens_name) for lens_name, lengths in focal_lengths.items()) 251 lens_names_by_focal_length = [item[1] for item in lens_names_by_focal_length] 252 with open("lenses.txt", "w") as outfile: 253 if focal_lengths: 254 outfile.write("""# For suggestions for <maker> and <mount> see 255# <https://github.com/lensfun/lensfun/tree/master/data/db>. 256# <aspect-ratio> is optional and by default 3:2. 257# Omit <type> for rectilinear lenses. 258""") 259 for lens_name in lens_names_by_focal_length: 260 outfile.write("\n{0}: <maker>, <mount>, <cropfactor>, <aspect-ratio>, <type>\n".format(lens_name)) 261 # FixMe: Only generate focal lengths that are available for 262 # *distortion*. 263 for length in sorted(focal_lengths[lens_name]): 264 outfile.write("distortion({0}mm) = , , \n".format(length)) 265 else: 266 outfile.write("""# No RAW images found (or no focal lengths in them). 267# Please have a look at 268# http://wilson.bronger.org/lens_calibration_tutorial/ 269""") 270 print("I wrote a template lenses.txt. Please fill this file with proper information. Abort.") 271 sys.exit() 272 273 274# 275# TCA correction 276# 277 278def generate_tca_tiffs(filename): 279 tca_filename = filename + ".tca" 280 if not os.path.exists(tca_filename): 281 tiff_filename = os.path.splitext(filename)[0] + ".tiff" 282 if not os.path.exists(tiff_filename): 283 subprocess.check_call(generate_raw_conversion_call(filename, ["-4", "-o", "0", "-M"])) 284 return filename, tiff_filename, tca_filename 285 return None, None, None 286 287if os.path.exists("tca"): 288 with chdir("tca"): 289 pool = multiprocessing.Pool() 290 results = set() 291 raw_files = find_raw_files() 292 for filename, tiff_filename, tca_filename in pool.map(generate_tca_tiffs, raw_files): 293 if filename: 294 output = subprocess.check_output(["tca_correct", "-o", "bv" if args.complex_tca else "v", tiff_filename], 295 stderr=open(os.devnull, "w")).splitlines()[-1].strip() 296 exif_data = file_exif_data[os.path.join("tca", filename)] 297 with open(tca_filename, "w") as outfile: 298 outfile.write("{0}\n{1}\n{2}\n".format(exif_data[0], exif_data[1], output.decode("ascii"))) 299 pool.close() 300 pool.join() 301 calibration_lines = {} 302 for filename in find_raw_files(): 303 lens_name, focal_length, tca_output = [line.rstrip("\n") for line in open(filename + ".tca").readlines()] 304 focal_length = float(focal_length) 305 data = re.match( 306 r"-r [.0]+:(?P<br>[-.0-9]+):[.0]+:(?P<vr>[-.0-9]+) -b [.0]+:(?P<bb>[-.0-9]+):[.0]+:(?P<vb>[-.0-9]+)", 307 tca_output).groupdict() 308 open(filename + "_tca.gp", "w").write( 309 """set title "{}" 310plot [0:1.8] {} * x**2 + {} title "red", {} * x**2 + {} title "blue" 311pause -1""".format(filename, data["br"], data["vr"], data["bb"], data["vb"])) 312 calibration_lines.setdefault(lens_name, []).append((focal_length, 313 """<tca model="poly3" focal="{0:g}" br="{1}" vr="{2}" bb="{3}" vb="{4}"/>""".format( 314 focal_length, data["br"], data["vr"], data["bb"], data["vb"]) if args.complex_tca else 315 """<tca model="poly3" focal="{0:g}" vr="{1}" vb="{2}"/>""".format( 316 focal_length, data["vr"], data["vb"]))) 317 for lens_name, lines in calibration_lines.items(): 318 lines.sort() 319 lenses[lens_name].calibration_lines.extend(line[1] for line in lines) 320 321 322# 323# Vignetting 324# 325 326images = {} 327distances_per_triplett = {} 328 329for vignetting_directory in glob.glob("vignetting*"): 330 distance = float(vignetting_directory.partition("_")[2] or "inf") 331 assert distance == float("inf") or distance < 1000 332 with chdir(vignetting_directory): 333 for filename in find_raw_files(): 334 exif_data = file_exif_data[os.path.join(vignetting_directory, filename)] + (distance,) 335 distances_per_triplett.setdefault(exif_data[:3], set()).add(distance) 336 images.setdefault(exif_data, []).append(os.path.join(vignetting_directory, filename)) 337 338vignetting_db_entries = {} 339 340def evaluate_image_set(exif_data, filepaths): 341 output_filename = "{0}--{1}--{2}--{3}".format(*exif_data).replace(" ", "_").replace("/", "__").replace(":", "___"). \ 342 replace("*", "++").replace("=", "##") 343 gnuplot_filename = "{0}.gp".format(output_filename) 344 try: 345 gnuplot_line = codecs.open(gnuplot_filename, encoding="utf-8").readlines()[3] 346 match = re.match(r' [-e.0-9]+ \* \(1 \+ \((?P<k1>[-e.0-9]+)\) \* x\*\*2 \+ \((?P<k2>[-e.0-9]+)\) \* x\*\*4 \+ ' 347 r'\((?P<k3>[-e.0-9]+)\) \* x\*\*6\) title "fit"', gnuplot_line) 348 k1, k2, k3 = [float(k) for k in match.groups()] 349 except IOError: 350 radii, intensities = [], [] 351 for filepath in filepaths: 352 maximal_radius = 1 353 try: 354 sidecar_file = open(os.path.splitext(filepath)[0] + ".txt") 355 except FileNotFoundError: 356 pass 357 else: 358 for line in sidecar_file: 359 if line.startswith("maximal_radius"): 360 maximal_radius = float(line.partition(":")[2]) 361 dcraw_process = subprocess.Popen(generate_raw_conversion_call(filepath, ["-4", "-M", "-o", "0", "-c"] + h_option), 362 stdout=subprocess.PIPE) 363 image_data = subprocess.check_output( 364 ["convert", "tiff:-", "-set", "colorspace", "RGB", "-resize", "250", "pgm:-"], stdin=dcraw_process.stdout, 365 stderr=open(os.devnull, "w")) 366 width, height = None, None 367 header_size = 0 368 for i, line in enumerate(image_data.splitlines(True)): 369 header_size += len(line) 370 if i == 0: 371 assert line == b"P5\n", "Wrong image format (must be NetPGM binary)" 372 else: 373 line = line.partition(b"#")[0].strip() 374 if line: 375 if not width: 376 width, height = line.split() 377 width, height = int(width), int(height) 378 else: 379 assert line == b"65535", "Wrong grayscale depth: {} (must be 65535)".format(int(line)) 380 break 381 half_diagonal = math.hypot(width // 2, height // 2) 382 image_data = struct.unpack("!{0}s{1}H".format(header_size, width * height), image_data)[1:] 383 for i, intensity in enumerate(image_data): 384 y, x = divmod(i, width) 385 radius = math.hypot(x - width // 2, y - height // 2) / half_diagonal 386 if radius <= maximal_radius: 387 radii.append(radius) 388 intensities.append(intensity) 389 all_points_filename = "{0}-all_points.dat".format(output_filename) 390 with open(all_points_filename, "w") as outfile: 391 for radius, intensity in zip(radii, intensities): 392 outfile.write("{0} {1}\n".format(radius, intensity)) 393 number_of_bins = 16 394 bins = [[] for i in range(number_of_bins)] 395 for radius, intensity in zip(radii, intensities): 396 # The zeroth and the last bin are only half bins which means that their 397 # means are skewed. But this is okay: For the zeroth, the curve is 398 # supposed to be horizontal anyway, and for the last, it underestimates 399 # the vignetting at the rim which is a good thing (too much of 400 # correction is bad). 401 bin_index = int(round(radius / maximal_radius * (number_of_bins - 1))) 402 bins[bin_index].append(intensity) 403 radii = [i / (number_of_bins - 1) * maximal_radius for i in range(number_of_bins)] 404 intensities = [numpy.median(bin) for bin in bins] 405 bins_filename = "{0}-bins.dat".format(output_filename) 406 with open(bins_filename, "w") as outfile: 407 for radius, intensity in zip(radii, intensities): 408 outfile.write("{0} {1}\n".format(radius, intensity)) 409 radii, intensities = numpy.array(radii), numpy.array(intensities) 410 411 def fit_function(radius, A, k1, k2, k3): 412 return A * (1 + k1 * radius**2 + k2 * radius**4 + k3 * radius**6) 413 414 A, k1, k2, k3 = leastsq(lambda p, x, y: y - fit_function(x, *p), [30000, -0.3, 0, 0], args=(radii, intensities))[0] 415 416 lens_name, focal_length, aperture, distance = exif_data 417 if distance == float("inf"): 418 distance = "∞" 419 codecs.open(gnuplot_filename, "w", encoding="utf-8").write("""set grid 420set title "{6}, {7} mm, f/{8}, {9} m" 421plot "{0}" with dots title "samples", "{1}" with linespoints lw 4 title "average", \\ 422 {2} * (1 + ({3}) * x**2 + ({4}) * x**4 + ({5}) * x**6) title "fit" 423pause -1 424""".format(all_points_filename, bins_filename, A, k1, k2, k3, lens_name, focal_length, aperture, distance)) 425 426 return (k1, k2, k3) 427 428pool = multiprocessing.Pool() 429results = {} 430for exif_data, filepaths in images.items(): 431 results[exif_data] = pool.apply_async(evaluate_image_set, [exif_data, filepaths]) 432for exif_data, result in results.items(): 433 vignetting_db_entries[exif_data] = result.get() 434pool.close() 435pool.join() 436 437new_vignetting_db_entries = {} 438for configuration, vignetting in vignetting_db_entries.items(): 439 triplett = configuration[:3] 440 distances = distances_per_triplett[triplett] 441 if len(distances) == 1 and list(distances)[0] > 10: 442 # If only one distance was measured and at more than 10m, insert it twice 443 # at 10m and ∞, so that the whole range is covered. 444 new_vignetting_db_entries[tuple(configuration[:3] + (10,))] = \ 445 new_vignetting_db_entries[tuple(configuration[:3] + (float("inf"),))] = vignetting 446 else: 447 new_vignetting_db_entries[configuration] = vignetting 448vignetting_db_entries = new_vignetting_db_entries 449 450 451for configuration in sorted(vignetting_db_entries): 452 lens, focal_length, aperture, distance = configuration 453 vignetting = vignetting_db_entries[configuration] 454 if distance == float("inf"): 455 distance = 1000 456 try: 457 lenses[lens].calibration_lines.append( 458 """<vignetting model="pa" focal="{focal_length:g}" aperture="{aperture:g}" distance="{distance:g}" """ 459 """k1="{vignetting[0]:.4f}" k2="{vignetting[1]:.4f}" k3="{vignetting[2]:.4f}"/>""".format( 460 focal_length=focal_length, aperture=aperture, vignetting=vignetting, distance=distance)) 461 except KeyError: 462 print("""Lens "{0}" not found in lenses.txt. Abort.""".format(lens)) 463 sys.exit() 464 465 466outfile = open("lensfun.xml", "w") 467outfile.write("<lensdatabase>\n") 468for lens in sorted(lenses.values()): 469 lens.write(outfile) 470outfile.write("\n</lensdatabase>\n") 471