set strict as default typing mode in pytorch

Reviewed By: azad-meta

Differential Revision: D52957335

fbshipit-source-id: 1f45063e5c6e15d8ae3305efeef3b492131a858a

torchbenchmark/models/Background_Matting/Data_adobe/compose.py

@@ -1,134 +1,134 @@
-##Copyright 2017 Adobe Systems Inc.
-##
-##Licensed under the Apache License, Version 2.0 (the "License");
-##you may not use this file except in compliance with the License.
-##You may obtain a copy of the License at
-##
-##    http://www.apache.org/licenses/LICENSE-2.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.
-
-
-##############################################################
-# python compose.py --fg_path fg_train --mask_path mask_train --bg_path bg_train --out_path merged_train --out_csv Adobe_train_data.csv --workers 8
-
-from PIL import Image
-from tqdm import tqdm
-import argparse
-import os
-import logging
-import math
-from multiprocessing.pool import ThreadPool
-import threading
-
-parser = argparse.ArgumentParser(description='compose backgrounds and foregrounds')
-
-parser.add_argument('--fg_path', type=str, required=True, help='path to provided foreground images')
-parser.add_argument('--mask_path', type=str, required=True, help='path to provided alpha mattes')
-parser.add_argument('--bg_path', type=str, required=True, help='path to to background images (MSCOCO)')
-parser.add_argument('--out_path', type=str, required=True, help='path to folder where you want the composited images to go')
-
-parser.add_argument('--out_csv', type=str, default=os.devnull, help='path to csv file used by data loader')
-parser.add_argument('--num_bgs', type=int, default=100, help='number of backgrounds onto which to paste each foreground')
-parser.add_argument('--workers', type=int, default=1, help='maximum workers to use, defaults to 1')
-args = parser.parse_args()
-
-fg_path, a_path, bg_path, out_path, num_bgs = args.fg_path, args.mask_path, args.bg_path, args.out_path, args.num_bgs
-os.makedirs(out_path, exist_ok=True)
-
-def format_pbar_str(i, im_name):
-    pbar_prefix = "(" + str(i) + ") "
-    width = 33 - len(pbar_prefix)
-    pretty_name = pbar_prefix + ("..." + im_name[-(width - 3):] if len(im_name) > width else im_name)
-    return pretty_name.rjust(33)
-
-def fixpath(path):
-    return 'Data_adobe/' + path if not os.path.isabs(path) else path
-
-def composite4(fg, bg, a, w, h):
-    bg = bg.crop((0,0,w,h))
-    bg.paste(fg, mask=a)
-    return bg
-
-def process_foreground_image(i, job):
-    worker_thread_id = int(threading.current_thread().name.rpartition("-")[-1])
-    im_name, bg_batch = job
-
-    im_name = im_name.replace(fg_path, '')
-    im = Image.open(os.path.join(fg_path, im_name))
-    al = Image.open(os.path.join(a_path, im_name))
-    bbox = im.size
-    w = bbox[0]
-    h = bbox[1]
-    if im.mode != 'RGB' and im.mode != 'RGBA':
-        im = im.convert('RGB')
-    if len(al.getbands()) > 0: # take the first channel, usually R
-        al = al.split()[0]
-
-    output_lines = []
-    with lock:
-        pbar = tqdm(bg_batch, position=worker_thread_id, desc=format_pbar_str(i, im_name), leave=False)
-    for b, bg_name in enumerate(pbar):
-        bg = Image.open(os.path.join(bg_path, bg_name))
-        if bg.mode != 'RGB':
-            bg = bg.convert('RGB')
-
-        bg_bbox = bg.size
-        bw = bg_bbox[0]
-        bh = bg_bbox[1]
-        wratio = w / bw
-        hratio = h / bh
-        ratio = wratio if wratio > hratio else hratio
-        if ratio > 1:
-            bg = bg.resize((math.ceil(bw * ratio), math.ceil(bh * ratio)), Image.BICUBIC)
-
-        try:
-            out = composite4(im, bg, al, w, h)
-            back_idx = i * num_bgs + b
-            out_name = os.path.join(out_path, im_name[:len(im_name) - 4] + '_' + str(back_idx) + '_comp.png')
-            out.save(out_name, "PNG")
-
-            back = bg.crop((0, 0, w, h))
-            back_name = os.path.join(out_path, im_name[:len(im_name) - 4] + '_' + str(back_idx) + '_back.png')
-            back.save(back_name, "PNG")
-
-            line = os.path.join(fixpath(fg_path), im_name) + ';' + os.path.join(fixpath(a_path), im_name) + ';' + fixpath(out_name) + ';' + fixpath(back_name) + '\n'
-            output_lines.append(line)
-        except Exception as e:
-            logging.error(f"Composing {im_name} onto {bg_name} failed! Skipping. Error: %s" % e)
-        with lock:
-            pbar.update()
-    with lock:
-        pbar.close()
-    return output_lines
-
-fg_files = os.listdir(fg_path)
-a_files = os.listdir(a_path)
-bg_files = os.listdir(bg_path)
-bg_batches = [bg_files[i * num_bgs:(i + 1) * num_bgs] for i in range((len(bg_files) + num_bgs - 1) // num_bgs )]
-
-lock = threading.Lock()
-pool = ThreadPool(args.workers)
-with lock:
-    total_pbar = tqdm(total=len(fg_files), position=args.workers+2, desc="TOTAL", leave=True, smoothing=0.0)
-def update_total_pbar(_):
-    with lock:
-        total_pbar.update(1)
-jobs = []
-for jobargs in enumerate(zip(fg_files, bg_batches)):
-    jobs.append(pool.apply_async(process_foreground_image, args=jobargs, callback=update_total_pbar))
-pool.close()
-pool.join()
-
-output = []
-for result in jobs:
-    output.extend(result.get())
-tqdm.write("Done composing...")
-
-with open(args.out_csv, "w") as f:
-    for line in output:
-        f.write(line)
+##Copyright 2017 Adobe Systems Inc.
+##
+##Licensed under the Apache License, Version 2.0 (the "License");
+##you may not use this file except in compliance with the License.
+##You may obtain a copy of the License at
+##
+##    http://www.apache.org/licenses/LICENSE-2.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.
+
+
+##############################################################
+# python compose.py --fg_path fg_train --mask_path mask_train --bg_path bg_train --out_path merged_train --out_csv Adobe_train_data.csv --workers 8
+
+from PIL import Image
+from tqdm import tqdm
+import argparse
+import os
+import logging
+import math
+from multiprocessing.pool import ThreadPool
+import threading
+
+parser = argparse.ArgumentParser(description='compose backgrounds and foregrounds')
+
+parser.add_argument('--fg_path', type=str, required=True, help='path to provided foreground images')
+parser.add_argument('--mask_path', type=str, required=True, help='path to provided alpha mattes')
+parser.add_argument('--bg_path', type=str, required=True, help='path to to background images (MSCOCO)')
+parser.add_argument('--out_path', type=str, required=True, help='path to folder where you want the composited images to go')
+
+parser.add_argument('--out_csv', type=str, default=os.devnull, help='path to csv file used by data loader')
+parser.add_argument('--num_bgs', type=int, default=100, help='number of backgrounds onto which to paste each foreground')
+parser.add_argument('--workers', type=int, default=1, help='maximum workers to use, defaults to 1')
+args = parser.parse_args()
+
+fg_path, a_path, bg_path, out_path, num_bgs = args.fg_path, args.mask_path, args.bg_path, args.out_path, args.num_bgs
+os.makedirs(out_path, exist_ok=True)
+
+def format_pbar_str(i, im_name):
+    pbar_prefix = "(" + str(i) + ") "
+    width = 33 - len(pbar_prefix)
+    pretty_name = pbar_prefix + ("..." + im_name[-(width - 3):] if len(im_name) > width else im_name)
+    return pretty_name.rjust(33)
+
+def fixpath(path):
+    return 'Data_adobe/' + path if not os.path.isabs(path) else path
+
+def composite4(fg, bg, a, w, h):
+    bg = bg.crop((0,0,w,h))
+    bg.paste(fg, mask=a)
+    return bg
+
+def process_foreground_image(i, job):
+    worker_thread_id = int(threading.current_thread().name.rpartition("-")[-1])
+    im_name, bg_batch = job
+
+    im_name = im_name.replace(fg_path, '')
+    im = Image.open(os.path.join(fg_path, im_name))
+    al = Image.open(os.path.join(a_path, im_name))
+    bbox = im.size
+    w = bbox[0]
+    h = bbox[1]
+    if im.mode != 'RGB' and im.mode != 'RGBA':
+        im = im.convert('RGB')
+    if len(al.getbands()) > 0: # take the first channel, usually R
+        al = al.split()[0]
+
+    output_lines = []
+    with lock:
+        pbar = tqdm(bg_batch, position=worker_thread_id, desc=format_pbar_str(i, im_name), leave=False)
+    for b, bg_name in enumerate(pbar):
+        bg = Image.open(os.path.join(bg_path, bg_name))
+        if bg.mode != 'RGB':
+            bg = bg.convert('RGB')
+
+        bg_bbox = bg.size
+        bw = bg_bbox[0]
+        bh = bg_bbox[1]
+        wratio = w / bw
+        hratio = h / bh
+        ratio = wratio if wratio > hratio else hratio
+        if ratio > 1:
+            bg = bg.resize((math.ceil(bw * ratio), math.ceil(bh * ratio)), Image.BICUBIC)
+
+        try:
+            out = composite4(im, bg, al, w, h)
+            back_idx = i * num_bgs + b
+            out_name = os.path.join(out_path, im_name[:len(im_name) - 4] + '_' + str(back_idx) + '_comp.png')
+            out.save(out_name, "PNG")
+
+            back = bg.crop((0, 0, w, h))
+            back_name = os.path.join(out_path, im_name[:len(im_name) - 4] + '_' + str(back_idx) + '_back.png')
+            back.save(back_name, "PNG")
+
+            line = os.path.join(fixpath(fg_path), im_name) + ';' + os.path.join(fixpath(a_path), im_name) + ';' + fixpath(out_name) + ';' + fixpath(back_name) + '\n'
+            output_lines.append(line)
+        except Exception as e:
+            logging.error(f"Composing {im_name} onto {bg_name} failed! Skipping. Error: %s" % e)
+        with lock:
+            pbar.update()
+    with lock:
+        pbar.close()
+    return output_lines
+
+fg_files = os.listdir(fg_path)
+a_files = os.listdir(a_path)
+bg_files = os.listdir(bg_path)
+bg_batches = [bg_files[i * num_bgs:(i + 1) * num_bgs] for i in range((len(bg_files) + num_bgs - 1) // num_bgs )]
+
+lock = threading.Lock()
+pool = ThreadPool(args.workers)
+with lock:
+    total_pbar = tqdm(total=len(fg_files), position=args.workers+2, desc="TOTAL", leave=True, smoothing=0.0)
+def update_total_pbar(_):
+    with lock:
+        total_pbar.update(1)
+jobs = []
+for jobargs in enumerate(zip(fg_files, bg_batches)):
+    jobs.append(pool.apply_async(process_foreground_image, args=jobargs, callback=update_total_pbar))
+pool.close()
+pool.join()
+
+output = []
+for result in jobs:
+    output.extend(result.get())
+tqdm.write("Done composing...")
+
+with open(args.out_csv, "w") as f:
+    for line in output:
+        f.write(line)

torchbenchmark/models/pytorch_unet/pytorch_unet/utils/utils.py

@@ -1,17 +1,17 @@
-import matplotlib.pyplot as plt
-
-
-def plot_img_and_mask(img, mask):
-    classes = mask.shape[0] if len(mask.shape) > 2 else 1
-    fig, ax = plt.subplots(1, classes + 1)
-    ax[0].set_title('Input image')
-    ax[0].imshow(img)
-    if classes > 1:
-        for i in range(classes):
-            ax[i + 1].set_title(f'Output mask (class {i + 1})')
-            ax[i + 1].imshow(mask[:, :, i])
-    else:
-        ax[1].set_title(f'Output mask')
-        ax[1].imshow(mask)
-    plt.xticks([]), plt.yticks([])
-    plt.show()
+import matplotlib.pyplot as plt
+
+
+def plot_img_and_mask(img, mask):
+    classes = mask.shape[0] if len(mask.shape) > 2 else 1
+    fig, ax = plt.subplots(1, classes + 1)
+    ax[0].set_title('Input image')
+    ax[0].imshow(img)
+    if classes > 1:
+        for i in range(classes):
+            ax[i + 1].set_title(f'Output mask (class {i + 1})')
+            ax[i + 1].imshow(mask[:, :, i])
+    else:
+        ax[1].set_title(f'Output mask')
+        ax[1].imshow(mask)
+    plt.xticks([]), plt.yticks([])
+    plt.show()