inpainting test

.aidigestignore

@@ -0,0 +1,10 @@
+controlnet_aux_local/normalbae/*
+requirements.txt
+win.requirements.txt
+web.html
+client.py
+local_app.py
+README.md
+Dockerfile
+.gitignore
+.gitattributes

app.py

@@ -15,11 +15,13 @@ import imageio
 from huggingface_hub import HfApi
 import gc
 import torch
+import cv2
 from PIL import Image
 from diffusers import (
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
+    StableDiffusionInpaintPipeline,
     # AutoencoderKL,
 )
 from controlnet_aux_local import NormalBaeDetector
@@ -98,6 +100,19 @@ if gr.NO_RELOAD:
         # vae=vae,
         torch_dtype=torch.float16,
     ).to("cuda")
+    
+    print('loading inpainting pipe')
+    inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-inpainting",
+        torch_dtype=torch.float16,
+    ).to("cuda")
+    
+    print('loading controlnet inpainting pipe')
+    controlnet_inpaint_pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-inpainting",
+        controlnet=controlnet,
+        torch_dtype=torch.float16,
+    ).to("cuda")
 
     print("loading preprocessor")
     preprocessor = Preprocessor()
@@ -119,6 +134,21 @@ if gr.NO_RELOAD:
     gc.collect()
     print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")
     print("Model Compiled!")
+    
+def generate_furniture_mask(image, furniture_type):
+    image_np = np.array(image)
+    height, width = image_np.shape[:2]
+        
+    mask = np.zeros((height, width), dtype=np.uint8)
+        
+    if furniture_type == "sofa":
+        cv2.rectangle(mask, (width//4, int(height*0.6)), (width*3//4, height), 255, -1)
+    elif furniture_type == "table":
+        cv2.rectangle(mask, (width//3, height//3), (width*2//3, height*2//3), 255, -1)
+    elif furniture_type == "chair":
+        cv2.circle(mask, (width*3//5, height*2//3), height//6, 255, -1)
+        
+    return Image.fromarray(mask)
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -422,15 +452,46 @@ def process_image(
     print(prompt)
     print(f"\n-------------------------Preprocess done in: {preprocess_time:.2f} seconds-------------------------")    
     start = time.time()
-    results = pipe(
+    # results = pipe(
+    #     prompt=prompt,
+    #     negative_prompt=negative_prompt,
+    #     guidance_scale=guidance_scale,
+    #     num_images_per_prompt=num_images,
+    #     num_inference_steps=num_steps,
+    #     generator=generator,
+    #     image=control_image,
+    # ).images[0]
+    
+    initial_result = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
-        num_images_per_prompt=num_images,
+        num_images_per_prompt=1,
         num_inference_steps=num_steps,
         generator=generator,
         image=control_image,
     ).images[0]
+    
+    # Randomly choose whether to add furniture and which type
+    furniture_types = ["None", "sofa", "table", "chair"]
+    furniture_type = random.choice(furniture_types)
+    
+    if furniture_type != "None":
+        furniture_mask = generate_furniture_mask(initial_result, furniture_type)
+        furniture_prompt = f"A {furniture_type} in the style of {style_selection}"
+        inpainted_image = controlnet_inpaint_pipe(
+            prompt=furniture_prompt,
+            image=initial_result,
+            mask_image=furniture_mask,
+            control_image=control_image,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        ).images[0]
+    else:
+        inpainted_image = initial_result
+    
     print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
     torch.cuda.empty_cache()
     
@@ -456,7 +517,7 @@ def process_image(
         token=API_KEY,
         run_as_future=True,
     )
-    return results
+    return inpainted_image
 
 if prod:
     demo.queue(max_size=20).launch(server_name="localhost", server_port=port)

codebase.md

@@ -0,0 +1,843 @@
+# preprocess.py
+
+```py
+import PIL.Image
+import torch, gc
+from controlnet_aux_local import NormalBaeDetector#, CannyDetector
+
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
+
+    def __init__(self):
+        self.model = None
+        self.name = ""
+
+    def load(self, name: str) -> None:
+        if name == self.name:
+            return
+        elif name == "NormalBae":
+            print("Loading NormalBae")
+            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
+            torch.cuda.empty_cache()
+            self.name = name
+        else:
+            raise ValueError
+        return
+
+    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
+        return self.model(image, **kwargs)
+```
+
+# app.py
+
+```py
+prod = False
+port = 8080
+show_options = False
+if prod:
+    port = 8081
+    # show_options = False
+
+import os
+import random
+import time
+import gradio as gr
+import numpy as np
+import spaces
+import imageio
+from huggingface_hub import HfApi
+import gc
+import torch
+from PIL import Image
+from diffusers import (
+    ControlNetModel,
+    DPMSolverMultistepScheduler,
+    StableDiffusionControlNetPipeline,
+    # AutoencoderKL,
+)
+from controlnet_aux_local import NormalBaeDetector
+
+MAX_SEED = np.iinfo(np.int32).max
+API_KEY = os.environ.get("API_KEY", None)
+# os.environ['HF_HOME'] = '/data/.huggingface'
+
+print("CUDA version:", torch.version.cuda)
+print("loading everything")
+compiled = False
+api = HfApi()
+
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
+
+    def __init__(self):
+        self.model = None
+        self.name = ""
+
+    def load(self, name: str) -> None:
+        if name == self.name:
+            return
+        elif name == "NormalBae":
+            print("Loading NormalBae")
+            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
+            torch.cuda.empty_cache()
+            self.name = name
+        else:
+            raise ValueError
+        return
+
+    def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
+        return self.model(image, **kwargs)
+
+if gr.NO_RELOAD:
+    # Controlnet Normal
+    model_id = "lllyasviel/control_v11p_sd15_normalbae"
+    print("initializing controlnet")
+    controlnet = ControlNetModel.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        attn_implementation="flash_attention_2",
+    ).to("cuda")
+
+    # Scheduler
+    scheduler = DPMSolverMultistepScheduler.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        solver_order=2,
+        subfolder="scheduler",
+        use_karras_sigmas=True,
+        final_sigmas_type="sigma_min",
+        algorithm_type="sde-dpmsolver++",
+        prediction_type="epsilon",
+        thresholding=False,
+        denoise_final=True,
+        device_map="cuda",
+        torch_dtype=torch.float16,
+    )
+
+    # Stable Diffusion Pipeline URL
+    # base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
+    base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
+    # vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
+
+    # print('loading vae')
+    # vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
+    # vae.to(memory_format=torch.channels_last) 
+
+    print('loading pipe')
+    pipe = StableDiffusionControlNetPipeline.from_single_file(
+        base_model_url,
+        safety_checker=None,
+        controlnet=controlnet,
+        scheduler=scheduler,
+        # vae=vae,
+        torch_dtype=torch.float16,
+    ).to("cuda")
+
+    print("loading preprocessor")
+    preprocessor = Preprocessor()
+    preprocessor.load("NormalBae")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2",)
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_unconventional maid.pt", token="HDA_unconventional_maid")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress")
+    # pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari")
+    pipe.to("cuda")
+
+    print("---------------Loaded controlnet pipeline---------------") 
+    torch.cuda.empty_cache()
+    gc.collect()
+    print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")
+    print("Model Compiled!")
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+def get_additional_prompt():
+    prompt = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
+    top = ["tank top", "blouse", "button up shirt", "sweater", "corset top"]
+    bottom = ["short skirt", "athletic shorts", "jean shorts", "pleated skirt", "short skirt", "leggings", "high-waisted shorts"]
+    accessory = ["knee-high boots", "gloves", "Thigh-high stockings", "Garter belt", "choker", "necklace", "headband", "headphones"]
+    return f"{prompt}, {random.choice(top)}, {random.choice(bottom)}, {random.choice(accessory)}, score_9"
+    # outfit = ["schoolgirl outfit", "playboy outfit", "red dress", "gala dress", "cheerleader outfit", "nurse outfit", "Kimono"]
+
+def get_prompt(prompt, additional_prompt):
+    interior = "design-style interior designed (interior space),tungsten white balance,captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length"
+    default = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
+    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
+    randomize = get_additional_prompt()
+    # nude = "NSFW,((nude)),medium bare breasts,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
+    # bodypaint = "((fully naked with no clothes)),nude naked seethroughxray,invisiblebodypaint,rating_newd,NSFW"
+    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
+    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
+    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
+    # ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao"
+    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
+    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
+    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
+    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
+    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
+    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
+    abg = "(1girl, asian body covered in words, words on body, tattoos of (words) on body),(masterpiece, best quality),medium breasts,(intricate details),unity 8k wallpaper,ultra detailed,(pastel colors),beautiful and aesthetic,see-through (clothes),detailed,solo"
+    # shibari = "extremely detailed, hyperrealistic photography, earrings, blushing, lace choker, tattoo, medium hair, score_9, HDA_Shibari"
+    shibari2 = "octane render, highly detailed, volumetric, HDA_Shibari"
+    
+    if prompt == "":
+        girls = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari2, ahegao2]
+        prompts_nsfw = [abg, shibari2, ahegao2]
+        prompt = f"{random.choice(girls)}"
+        prompt = f"boho chic"
+        # print(f"-------------{preset}-------------")
+    else:
+        prompt = f"Photo from Pinterest of {prompt} {interior}"
+        # prompt = default2
+    return f"{prompt} f{additional_prompt}"
+
+style_list = [
+    {
+        "name": "None",
+        "prompt": ""
+    },
+    {
+        "name": "Minimalistic",
+        "prompt": "Minimalist interior design,clean lines,neutral colors,uncluttered space,functional furniture,lots of natural light"
+    },
+    {
+        "name": "Boho",
+        "prompt": "Bohemian chic interior,eclectic mix of patterns and textures,vintage furniture,plants,woven textiles,warm earthy colors"
+    },
+    {
+        "name": "Farmhouse",
+        "prompt": "Modern farmhouse interior,rustic wood elements,shiplap walls,neutral color palette,industrial accents,cozy textiles"
+    },
+    {
+        "name": "Saudi Prince",
+        "prompt": "Opulent gold interior,luxurious ornate furniture,crystal chandeliers,rich fabrics,marble floors,intricate Arabic patterns"
+    },
+    {
+        "name": "Neoclassical",
+        "prompt": "Neoclassical interior design,elegant columns,ornate moldings,symmetrical layout,refined furniture,muted color palette"
+    },
+    {
+        "name": "Eclectic",
+        "prompt": "Eclectic interior design,mix of styles and eras,bold color combinations,diverse furniture pieces,unique art objects"
+    },
+    {
+        "name": "Parisian",
+        "prompt": "Parisian apartment interior,all-white color scheme,ornate moldings,herringbone wood floors,elegant furniture,large windows"
+    },
+    {
+        "name": "Hollywood",
+        "prompt": "Hollywood Regency interior,glamorous and luxurious,bold colors,mirrored surfaces,velvet upholstery,gold accents"
+    },
+    {
+        "name": "Scandinavian",
+        "prompt": "Scandinavian interior design,light wood tones,white walls,minimalist furniture,cozy textiles,hygge atmosphere"
+    },
+    {
+        "name": "Beach",
+        "prompt": "Coastal beach house interior,light blue and white color scheme,weathered wood,nautical accents,sheer curtains,ocean view"
+    },
+    {
+        "name": "Japanese",
+        "prompt": "Traditional Japanese interior,tatami mats,shoji screens,low furniture,zen garden view,minimalist decor,natural materials"
+    },
+    { 
+        "name": "Midcentury Modern",
+        "prompt": "Mid-century modern interior,1950s-60s style furniture,organic shapes,warm wood tones,bold accent colors,large windows"
+    },
+    {
+        "name": "Retro Futurism",
+        "prompt": "Neon (atompunk world) retro cyberpunk background",
+    },
+    {
+        "name": "Texan",
+        "prompt": "Western cowboy interior,rustic wood beams,leather furniture,cowhide rugs,antler chandeliers,southwestern patterns"
+    },
+    {
+        "name": "Matrix",
+        "prompt": "Futuristic cyberpunk interior,neon accent lighting,holographic plants,sleek black surfaces,advanced gaming setup,transparent screens,Blade Runner inspired decor,high-tech minimalist furniture"
+    }
+] 
+
+styles = {k["name"]: (k["prompt"]) for k in style_list}
+STYLE_NAMES = list(styles.keys())
+
+def apply_style(style_name):
+    if style_name in styles:
+        p = styles.get(style_name, "none")
+    return p
+
+    
+css = """
+h1, h2, h3 {
+    text-align: center;
+    display: block;
+}
+footer {
+    visibility: hidden;
+}
+.gradio-container {
+    max-width: 1100px !important;
+}
+.gr-image {
+    display: flex;
+    justify-content: center; 
+    align-items: center;
+    width: 100%;
+    height: 512px;
+    overflow: hidden;
+}
+.gr-image img {
+    width: 100%;
+    height: 100%; 
+    object-fit: cover;
+    object-position: center;
+}
+"""
+with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
+    #############################################################################
+    with gr.Row():
+        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
+            num_images = gr.Slider(
+                label="Images", minimum=1, maximum=4, value=1, step=1
+            )
+            image_resolution = gr.Slider(
+                label="Image resolution",
+                minimum=256,
+                maximum=1024,
+                value=512,
+                step=256,
+            )
+            preprocess_resolution = gr.Slider(
+                label="Preprocess resolution",
+                minimum=128,
+                maximum=1024,
+                value=512,
+                step=1,
+            )
+            num_steps = gr.Slider(
+                label="Number of steps", minimum=1, maximum=100, value=15, step=1
+            )  # 20/4.5 or 12 without lora, 4 with lora
+            guidance_scale = gr.Slider(
+                label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
+            )  # 5 without lora, 2 with lora
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            a_prompt = gr.Textbox(
+                label="Additional prompt",
+                value = "design-style interior designed (interior space), tungsten white balance, captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length"
+            )
+            n_prompt = gr.Textbox(
+                label="Negative prompt",
+                value="EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)",
+            )
+    #############################################################################
+    # input text
+    with gr.Column():
+        prompt = gr.Textbox(
+            label="Custom Design",
+            placeholder="Enter a description (optional)",
+        )
+    # design options
+    with gr.Row(visible=True):
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value="None",
+            label="Design Styles",
+        )
+    # input image
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1, min_width=300):
+            image = gr.Image(
+                label="Input",
+                sources=["upload"],
+                show_label=True,
+                mirror_webcam=True,
+                type="pil",
+            )
+            # run button
+            with gr.Column():
+                run_button = gr.Button(value="Use this one", size="lg", visible=False)
+        # output image
+        with gr.Column(scale=1, min_width=300):
+            result = gr.Image(  
+                label="Output",
+                interactive=False,
+                type="pil",
+                show_share_button= False,
+            )
+            # Use this image button
+            with gr.Column():
+                use_ai_button = gr.Button(value="Use this one", size="lg", visible=False)
+    config = [
+        image,
+        style_selection,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_images,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        guidance_scale,
+        seed,
+    ]
+    
+    with gr.Row():
+        helper_text = gr.Markdown("## Tap and hold (on mobile) to save the image.", visible=True)
+    
+    # image processing
+    @gr.on(triggers=[image.upload, prompt.submit, run_button.click], inputs=config, outputs=result, show_progress="minimal")
+    def auto_process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
+        return process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
+    
+    # AI image processing
+    @gr.on(triggers=[use_ai_button.click], inputs=[result] + config, outputs=[image, result], show_progress="minimal")
+    def submit(previous_result, image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
+        # First, yield the previous result to update the input image immediately
+        yield previous_result, gr.update()
+        # Then, process the new input image
+        new_result = process_image(previous_result, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
+        # Finally, yield the new result
+        yield previous_result, new_result
+
+    # Turn off buttons when processing
+    @gr.on(triggers=[image.upload, use_ai_button.click, run_button.click], inputs=None, outputs=[run_button, use_ai_button], show_progress="hidden")
+    def turn_buttons_off():
+        return gr.update(visible=False), gr.update(visible=False)
+    
+    # Turn on buttons when processing is complete
+    @gr.on(triggers=[result.change], inputs=None, outputs=[use_ai_button, run_button], show_progress="hidden")
+    def turn_buttons_on():
+        return gr.update(visible=True), gr.update(visible=True)
+
+@spaces.GPU(duration=12)
+@torch.inference_mode()
+def process_image(
+    image,
+    style_selection,
+    prompt,
+    a_prompt,
+    n_prompt,
+    num_images,
+    image_resolution,
+    preprocess_resolution,
+    num_steps,
+    guidance_scale,
+    seed,
+):
+    preprocess_start = time.time()
+    print("processing image")
+
+    seed = random.randint(0, MAX_SEED)
+    generator = torch.cuda.manual_seed(seed)
+    preprocessor.load("NormalBae")
+    control_image = preprocessor(
+        image=image,
+        image_resolution=image_resolution,
+        detect_resolution=preprocess_resolution,
+    )
+    preprocess_time = time.time() - preprocess_start
+    if style_selection is not None or style_selection != "None":
+        prompt = "Photo from Pinterest of " + apply_style(style_selection) + " " + prompt + "," + a_prompt
+    else:
+        prompt=str(get_prompt(prompt, a_prompt))
+    negative_prompt=str(n_prompt)
+    print(prompt)
+    print(f"\n-------------------------Preprocess done in: {preprocess_time:.2f} seconds-------------------------")    
+    start = time.time()
+    results = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num_images,
+        num_inference_steps=num_steps,
+        generator=generator,
+        image=control_image,
+    ).images[0]
+    print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
+    torch.cuda.empty_cache()
+    
+    # upload block
+    timestamp = int(time.time())
+    img_path = f"{timestamp}.jpg"
+    results_path = f"{timestamp}_out.jpg"
+    imageio.imsave(img_path, image)
+    imageio.imsave(results_path, results)
+    api.upload_file(
+        path_or_fileobj=img_path,
+        path_in_repo=img_path,
+        repo_id="broyang/interior-ai-outputs",
+        repo_type="dataset",
+        token=API_KEY,
+        run_as_future=True,
+    )
+    api.upload_file(
+        path_or_fileobj=results_path,
+        path_in_repo=results_path,
+        repo_id="broyang/interior-ai-outputs",
+        repo_type="dataset",
+        token=API_KEY,
+        run_as_future=True,
+    )
+    return results
+
+if prod:
+    demo.queue(max_size=20).launch(server_name="localhost", server_port=port)
+else:
+    demo.queue(api_open=False).launch(show_api=False)
+```
+
+# .aidigestignore
+
+```
+controlnet_aux_local/normalbae/*
+requirements.txt
+win.requirements.txt
+web.html
+client.py
+local_app.py
+README.md
+Dockerfile
+.gitignore
+.gitattributes
+```
+
+# controlnet_aux_local/util.py
+
+```py
+import os
+import random
+
+import cv2
+import numpy as np
+import torch
+
+annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts')
+
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+
+def make_noise_disk(H, W, C, F):
+    noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
+    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
+    noise = noise[F: F + H, F: F + W]
+    noise -= np.min(noise)
+    noise /= np.max(noise)
+    if C == 1:
+        noise = noise[:, :, None]
+    return noise
+
+
+def nms(x, t, s):
+    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+
+    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+
+    y = np.zeros_like(x)
+
+    for f in [f1, f2, f3, f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z
+
+def min_max_norm(x):
+    x -= np.min(x)
+    x /= np.maximum(np.max(x), 1e-5)
+    return x
+
+
+def safe_step(x, step=2):
+    y = x.astype(np.float32) * float(step + 1)
+    y = y.astype(np.int32).astype(np.float32) / float(step)
+    return y
+
+
+def img2mask(img, H, W, low=10, high=90):
+    assert img.ndim == 3 or img.ndim == 2
+    assert img.dtype == np.uint8
+
+    if img.ndim == 3:
+        y = img[:, :, random.randrange(0, img.shape[2])]
+    else:
+        y = img
+
+    y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC)
+
+    if random.uniform(0, 1) < 0.5:
+        y = 255 - y
+
+    return y < np.percentile(y, random.randrange(low, high))
+
+
+def resize_image(input_image, resolution):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    H = int(np.round(H / 64.0)) * 64
+    W = int(np.round(W / 64.0)) * 64
+    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
+    return img
+
+
+def torch_gc():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
+def ade_palette():
+    """ADE20K palette that maps each class to RGB values."""
+    return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
+            [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
+            [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
+            [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
+            [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
+            [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
+            [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
+            [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
+            [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
+            [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
+            [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
+            [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
+            [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
+            [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
+            [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
+            [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
+            [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
+            [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
+            [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
+            [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
+            [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
+            [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
+            [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
+            [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
+            [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
+            [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
+            [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
+            [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
+            [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
+            [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
+            [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
+            [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
+            [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
+            [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
+            [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
+            [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
+            [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
+            [102, 255, 0], [92, 0, 255]]
+
+
+```
+
+# controlnet_aux_local/processor.py
+
+```py
+"""
+This file contains a Processor that can be used to process images with controlnet aux processors
+"""
+import io
+import logging
+from typing import Dict, Optional, Union
+
+from PIL import Image
+
+from controlnet_aux_local import (CannyDetector, ContentShuffleDetector, HEDdetector,
+                            LeresDetector, LineartAnimeDetector,
+                            LineartDetector, MediapipeFaceDetector,
+                            MidasDetector, MLSDdetector, NormalBaeDetector,
+                            OpenposeDetector, PidiNetDetector, ZoeDetector,
+                            DWposeDetector)
+
+LOGGER = logging.getLogger(__name__)
+
+
+MODELS = {
+    # checkpoint models
+    'scribble_hed': {'class': HEDdetector, 'checkpoint': True},
+    'softedge_hed': {'class': HEDdetector, 'checkpoint': True},
+    'scribble_hedsafe': {'class': HEDdetector, 'checkpoint': True},
+    'softedge_hedsafe': {'class': HEDdetector, 'checkpoint': True},
+    'depth_midas': {'class': MidasDetector, 'checkpoint': True},
+    'mlsd': {'class': MLSDdetector, 'checkpoint': True},
+    'openpose': {'class': OpenposeDetector, 'checkpoint': True},
+    'openpose_face': {'class': OpenposeDetector, 'checkpoint': True},
+    'openpose_faceonly': {'class': OpenposeDetector, 'checkpoint': True},
+    'openpose_full': {'class': OpenposeDetector, 'checkpoint': True},
+    'openpose_hand': {'class': OpenposeDetector, 'checkpoint': True},
+    'dwpose': {'class': DWposeDetector, 'checkpoint': True},
+    'scribble_pidinet': {'class': PidiNetDetector, 'checkpoint': True},
+    'softedge_pidinet': {'class': PidiNetDetector, 'checkpoint': True},
+    'scribble_pidsafe': {'class': PidiNetDetector, 'checkpoint': True},
+    'softedge_pidsafe': {'class': PidiNetDetector, 'checkpoint': True},
+    'normal_bae': {'class': NormalBaeDetector, 'checkpoint': True},
+    'lineart_coarse': {'class': LineartDetector, 'checkpoint': True},
+    'lineart_realistic': {'class': LineartDetector, 'checkpoint': True},
+    'lineart_anime': {'class': LineartAnimeDetector, 'checkpoint': True},
+    'depth_zoe': {'class': ZoeDetector, 'checkpoint': True}, 
+    'depth_leres': {'class': LeresDetector, 'checkpoint': True}, 
+    'depth_leres++': {'class': LeresDetector, 'checkpoint': True}, 
+    # instantiate
+    'shuffle': {'class': ContentShuffleDetector, 'checkpoint': False},
+    'mediapipe_face': {'class': MediapipeFaceDetector, 'checkpoint': False},
+    'canny': {'class': CannyDetector, 'checkpoint': False},
+}
+
+
+MODEL_PARAMS = {
+    'scribble_hed': {'scribble': True},
+    'softedge_hed': {'scribble': False},
+    'scribble_hedsafe': {'scribble': True, 'safe': True},
+    'softedge_hedsafe': {'scribble': False, 'safe': True},
+    'depth_midas': {},
+    'mlsd': {},
+    'openpose': {'include_body': True, 'include_hand': False, 'include_face': False},
+    'openpose_face': {'include_body': True, 'include_hand': False, 'include_face': True},
+    'openpose_faceonly': {'include_body': False, 'include_hand': False, 'include_face': True},
+    'openpose_full': {'include_body': True, 'include_hand': True, 'include_face': True},
+    'openpose_hand': {'include_body': False, 'include_hand': True, 'include_face': False},
+    'dwpose': {},
+    'scribble_pidinet': {'safe': False, 'scribble': True},
+    'softedge_pidinet': {'safe': False, 'scribble': False},
+    'scribble_pidsafe': {'safe': True, 'scribble': True},
+    'softedge_pidsafe': {'safe': True, 'scribble': False},
+    'normal_bae': {},
+    'lineart_realistic': {'coarse': False},
+    'lineart_coarse': {'coarse': True},
+    'lineart_anime': {},
+    'canny': {},
+    'shuffle': {},
+    'depth_zoe': {},
+    'depth_leres': {'boost': False},
+    'depth_leres++': {'boost': True},
+    'mediapipe_face': {},
+}
+
+CHOICES = f"Choices for the processor are {list(MODELS.keys())}"
+
+
+class Processor:
+    def __init__(self, processor_id: str, params: Optional[Dict] = None) -> None:
+        """Processor that can be used to process images with controlnet aux processors
+
+        Args:
+            processor_id (str): processor name, options are 'hed, midas, mlsd, openpose,
+                                pidinet, normalbae, lineart, lineart_coarse, lineart_anime,
+                                canny, content_shuffle, zoe, mediapipe_face
+            params (Optional[Dict]): parameters for the processor
+        """
+        LOGGER.info(f"Loading {processor_id}")
+
+        if processor_id not in MODELS:
+            raise ValueError(f"{processor_id} is not a valid processor id. Please make sure to choose one of {', '.join(MODELS.keys())}")
+
+        self.processor_id = processor_id
+        self.processor = self.load_processor(self.processor_id)
+
+        # load default params
+        self.params = MODEL_PARAMS[self.processor_id]
+        # update with user params
+        if params:
+            self.params.update(params)
+
+    def load_processor(self, processor_id: str) -> 'Processor':
+        """Load controlnet aux processors
+
+        Args:
+            processor_id (str): processor name
+
+        Returns:
+            Processor: controlnet aux processor
+        """
+        processor = MODELS[processor_id]['class']
+
+        # check if the proecssor is a checkpoint model
+        if MODELS[processor_id]['checkpoint']:
+            processor = processor.from_pretrained("lllyasviel/Annotators")
+        else:
+            processor = processor()
+        return processor
+
+    def __call__(self, image: Union[Image.Image, bytes],
+                 to_pil: bool = True) -> Union[Image.Image, bytes]:
+        """processes an image with a controlnet aux processor
+
+        Args:
+            image (Union[Image.Image, bytes]): input image in bytes or PIL Image
+            to_pil (bool): whether to return bytes or PIL Image
+
+        Returns:
+            Union[Image.Image, bytes]: processed image in bytes or PIL Image
+        """
+        # check if bytes or PIL Image
+        if isinstance(image, bytes):
+            image = Image.open(io.BytesIO(image)).convert("RGB")
+
+        processed_image = self.processor(image, **self.params)
+
+        if to_pil:
+            return processed_image
+        else:
+            output_bytes = io.BytesIO()
+            processed_image.save(output_bytes, format='JPEG')
+            return output_bytes.getvalue()
+
+```
+
+# controlnet_aux_local/__init__.py
+
+```py
+__version__ = "0.0.8"
+
+# from .hed import HEDdetector
+# from .leres import LeresDetector
+# from .lineart import LineartDetector
+# from .lineart_anime import LineartAnimeDetector
+# from .midas import MidasDetector
+# from .mlsd import MLSDdetector
+from .normalbae import NormalBaeDetector
+# from .open_pose import OpenposeDetector
+# from .pidi import PidiNetDetector
+# from .zoe import ZoeDetector
+
+# from .canny import CannyDetector
+# from .mediapipe_face import MediapipeFaceDetector
+# from .segment_anything import SamDetector
+# from .shuffle import ContentShuffleDetector
+# from .dwpose import DWposeDetector
+```
+

controlnet_aux_local/canny/__init__.py

@@ -1,36 +0,0 @@
-import warnings
-import cv2
-import numpy as np
-from PIL import Image
-from ..util import HWC3, resize_image
-
-class CannyDetector:
-    def __call__(self, input_image=None, low_threshold=100, high_threshold=200, detect_resolution=512, image_resolution=512, output_type=None, **kwargs):
-        if "img" in kwargs:
-            warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
-            input_image = kwargs.pop("img")
-        
-        if input_image is None:
-            raise ValueError("input_image must be defined.")
-
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-            output_type = output_type or "pil"
-        else:
-            output_type = output_type or "np"
-        
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-
-        detected_map = cv2.Canny(input_image, low_threshold, high_threshold)
-        detected_map = HWC3(detected_map)      
-         
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-        
-        if output_type == "pil":
-            detected_map = Image.fromarray(detected_map)
-            
-        return detected_map

controlnet_aux_local/dwpose/__init__.py

@@ -1,91 +0,0 @@
-# Openpose
-# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
-# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
-# 3rd Edited by ControlNet
-# 4th Edited by ControlNet (added face and correct hands)
-
-import os
-os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
-
-import cv2
-import torch
-import numpy as np
-from PIL import Image
-
-from ..util import HWC3, resize_image
-from . import util
-
-
-def draw_pose(pose, H, W):
-    bodies = pose['bodies']
-    faces = pose['faces']
-    hands = pose['hands']
-    candidate = bodies['candidate']
-    subset = bodies['subset']
-    
-    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
-    canvas = util.draw_bodypose(canvas, candidate, subset)
-    canvas = util.draw_handpose(canvas, hands)
-    canvas = util.draw_facepose(canvas, faces)
-
-    return canvas
-
-class DWposeDetector:
-    def __init__(self, det_config=None, det_ckpt=None, pose_config=None, pose_ckpt=None, device="cpu"):
-        from .wholebody import Wholebody
-
-        self.pose_estimation = Wholebody(det_config, det_ckpt, pose_config, pose_ckpt, device)
-    
-    def to(self, device):
-        self.pose_estimation.to(device)
-        return self
-    
-    def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
-        
-        input_image = cv2.cvtColor(np.array(input_image, dtype=np.uint8), cv2.COLOR_RGB2BGR)
-
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-        H, W, C = input_image.shape
-        
-        with torch.no_grad():
-            candidate, subset = self.pose_estimation(input_image)
-            nums, keys, locs = candidate.shape
-            candidate[..., 0] /= float(W)
-            candidate[..., 1] /= float(H)
-            body = candidate[:,:18].copy()
-            body = body.reshape(nums*18, locs)
-            score = subset[:,:18]
-            
-            for i in range(len(score)):
-                for j in range(len(score[i])):
-                    if score[i][j] > 0.3:
-                        score[i][j] = int(18*i+j)
-                    else:
-                        score[i][j] = -1
-
-            un_visible = subset<0.3
-            candidate[un_visible] = -1
-
-            foot = candidate[:,18:24]
-
-            faces = candidate[:,24:92]
-
-            hands = candidate[:,92:113]
-            hands = np.vstack([hands, candidate[:,113:]])
-            
-            bodies = dict(candidate=body, subset=score)
-            pose = dict(bodies=bodies, hands=hands, faces=faces)
-            
-            detected_map = draw_pose(pose, H, W)
-            detected_map = HWC3(detected_map)
-            
-            img = resize_image(input_image, image_resolution)
-            H, W, C = img.shape
-
-            detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-
-            if output_type == "pil":
-                detected_map = Image.fromarray(detected_map)
-                
-            return detected_map

controlnet_aux_local/dwpose/util.py

@@ -1,303 +0,0 @@
-import math
-import numpy as np
-import cv2
-
-
-eps = 0.01
-
-
-def smart_resize(x, s):
-    Ht, Wt = s
-    if x.ndim == 2:
-        Ho, Wo = x.shape
-        Co = 1
-    else:
-        Ho, Wo, Co = x.shape
-    if Co == 3 or Co == 1:
-        k = float(Ht + Wt) / float(Ho + Wo)
-        return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
-    else:
-        return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2)
-
-
-def smart_resize_k(x, fx, fy):
-    if x.ndim == 2:
-        Ho, Wo = x.shape
-        Co = 1
-    else:
-        Ho, Wo, Co = x.shape
-    Ht, Wt = Ho * fy, Wo * fx
-    if Co == 3 or Co == 1:
-        k = float(Ht + Wt) / float(Ho + Wo)
-        return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
-    else:
-        return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2)
-
-
-def padRightDownCorner(img, stride, padValue):
-    h = img.shape[0]
-    w = img.shape[1]
-
-    pad = 4 * [None]
-    pad[0] = 0 # up
-    pad[1] = 0 # left
-    pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
-    pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
-
-    img_padded = img
-    pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
-    img_padded = np.concatenate((pad_up, img_padded), axis=0)
-    pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
-    img_padded = np.concatenate((pad_left, img_padded), axis=1)
-    pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
-    img_padded = np.concatenate((img_padded, pad_down), axis=0)
-    pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
-    img_padded = np.concatenate((img_padded, pad_right), axis=1)
-
-    return img_padded, pad
-
-
-def transfer(model, model_weights):
-    transfered_model_weights = {}
-    for weights_name in model.state_dict().keys():
-        transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
-    return transfered_model_weights
-
-
-def draw_bodypose(canvas, candidate, subset):
-    H, W, C = canvas.shape
-    candidate = np.array(candidate)
-    subset = np.array(subset)
-
-    stickwidth = 4
-
-    limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
-               [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
-               [1, 16], [16, 18], [3, 17], [6, 18]]
-
-    colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
-              [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
-              [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
-
-    for i in range(17):
-        for n in range(len(subset)):
-            index = subset[n][np.array(limbSeq[i]) - 1]
-            if -1 in index:
-                continue
-            Y = candidate[index.astype(int), 0] * float(W)
-            X = candidate[index.astype(int), 1] * float(H)
-            mX = np.mean(X)
-            mY = np.mean(Y)
-            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
-            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
-            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
-            cv2.fillConvexPoly(canvas, polygon, colors[i])
-
-    canvas = (canvas * 0.6).astype(np.uint8)
-
-    for i in range(18):
-        for n in range(len(subset)):
-            index = int(subset[n][i])
-            if index == -1:
-                continue
-            x, y = candidate[index][0:2]
-            x = int(x * W)
-            y = int(y * H)
-            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
-
-    return canvas
-
-
-def draw_handpose(canvas, all_hand_peaks):
-    import matplotlib
-    
-    H, W, C = canvas.shape
-
-    edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
-             [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
-    
-    # (person_number*2, 21, 2)
-    for i in range(len(all_hand_peaks)):
-        peaks = all_hand_peaks[i]
-        peaks = np.array(peaks)
-        
-        for ie, e in enumerate(edges):
-
-            x1, y1 = peaks[e[0]]
-            x2, y2 = peaks[e[1]]
-            
-            x1 = int(x1 * W)
-            y1 = int(y1 * H)
-            x2 = int(x2 * W)
-            y2 = int(y2 * H)
-            if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
-                cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255, thickness=2)
-
-        for _, keyponit in enumerate(peaks):
-            x, y = keyponit
-
-            x = int(x * W)
-            y = int(y * H)
-            if x > eps and y > eps:
-                cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
-    return canvas
-
-
-def draw_facepose(canvas, all_lmks):
-    H, W, C = canvas.shape
-    for lmks in all_lmks:
-        lmks = np.array(lmks)
-        for lmk in lmks:
-            x, y = lmk
-            x = int(x * W)
-            y = int(y * H)
-            if x > eps and y > eps:
-                cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
-    return canvas
-
-
-# detect hand according to body pose keypoints
-# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
-def handDetect(candidate, subset, oriImg):
-    # right hand: wrist 4, elbow 3, shoulder 2
-    # left hand: wrist 7, elbow 6, shoulder 5
-    ratioWristElbow = 0.33
-    detect_result = []
-    image_height, image_width = oriImg.shape[0:2]
-    for person in subset.astype(int):
-        # if any of three not detected
-        has_left = np.sum(person[[5, 6, 7]] == -1) == 0
-        has_right = np.sum(person[[2, 3, 4]] == -1) == 0
-        if not (has_left or has_right):
-            continue
-        hands = []
-        #left hand
-        if has_left:
-            left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
-            x1, y1 = candidate[left_shoulder_index][:2]
-            x2, y2 = candidate[left_elbow_index][:2]
-            x3, y3 = candidate[left_wrist_index][:2]
-            hands.append([x1, y1, x2, y2, x3, y3, True])
-        # right hand
-        if has_right:
-            right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]]
-            x1, y1 = candidate[right_shoulder_index][:2]
-            x2, y2 = candidate[right_elbow_index][:2]
-            x3, y3 = candidate[right_wrist_index][:2]
-            hands.append([x1, y1, x2, y2, x3, y3, False])
-
-        for x1, y1, x2, y2, x3, y3, is_left in hands:
-            # pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
-            # handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
-            # handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
-            # const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
-            # const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
-            # handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
-            x = x3 + ratioWristElbow * (x3 - x2)
-            y = y3 + ratioWristElbow * (y3 - y2)
-            distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
-            distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
-            width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
-            # x-y refers to the center --> offset to topLeft point
-            # handRectangle.x -= handRectangle.width / 2.f;
-            # handRectangle.y -= handRectangle.height / 2.f;
-            x -= width / 2
-            y -= width / 2  # width = height
-            # overflow the image
-            if x < 0: x = 0
-            if y < 0: y = 0
-            width1 = width
-            width2 = width
-            if x + width > image_width: width1 = image_width - x
-            if y + width > image_height: width2 = image_height - y
-            width = min(width1, width2)
-            # the max hand box value is 20 pixels
-            if width >= 20:
-                detect_result.append([int(x), int(y), int(width), is_left])
-
-    '''
-    return value: [[x, y, w, True if left hand else False]].
-    width=height since the network require squared input.
-    x, y is the coordinate of top left 
-    '''
-    return detect_result
-
-
-# Written by Lvmin
-def faceDetect(candidate, subset, oriImg):
-    # left right eye ear 14 15 16 17
-    detect_result = []
-    image_height, image_width = oriImg.shape[0:2]
-    for person in subset.astype(int):
-        has_head = person[0] > -1
-        if not has_head:
-            continue
-
-        has_left_eye = person[14] > -1
-        has_right_eye = person[15] > -1
-        has_left_ear = person[16] > -1
-        has_right_ear = person[17] > -1
-
-        if not (has_left_eye or has_right_eye or has_left_ear or has_right_ear):
-            continue
-
-        head, left_eye, right_eye, left_ear, right_ear = person[[0, 14, 15, 16, 17]]
-
-        width = 0.0
-        x0, y0 = candidate[head][:2]
-
-        if has_left_eye:
-            x1, y1 = candidate[left_eye][:2]
-            d = max(abs(x0 - x1), abs(y0 - y1))
-            width = max(width, d * 3.0)
-
-        if has_right_eye:
-            x1, y1 = candidate[right_eye][:2]
-            d = max(abs(x0 - x1), abs(y0 - y1))
-            width = max(width, d * 3.0)
-
-        if has_left_ear:
-            x1, y1 = candidate[left_ear][:2]
-            d = max(abs(x0 - x1), abs(y0 - y1))
-            width = max(width, d * 1.5)
-
-        if has_right_ear:
-            x1, y1 = candidate[right_ear][:2]
-            d = max(abs(x0 - x1), abs(y0 - y1))
-            width = max(width, d * 1.5)
-
-        x, y = x0, y0
-
-        x -= width
-        y -= width
-
-        if x < 0:
-            x = 0
-
-        if y < 0:
-            y = 0
-
-        width1 = width * 2
-        width2 = width * 2
-
-        if x + width > image_width:
-            width1 = image_width - x
-
-        if y + width > image_height:
-            width2 = image_height - y
-
-        width = min(width1, width2)
-
-        if width >= 20:
-            detect_result.append([int(x), int(y), int(width)])
-
-    return detect_result
-
-
-# get max index of 2d array
-def npmax(array):
-    arrayindex = array.argmax(1)
-    arrayvalue = array.max(1)
-    i = arrayvalue.argmax()
-    j = arrayindex[i]
-    return i, j

controlnet_aux_local/dwpose/wholebody.py

@@ -1,121 +0,0 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-import os
-import numpy as np
-import warnings
-
-try:
-    import mmcv
-except ImportError:
-    warnings.warn(
-        "The module 'mmcv' is not installed. The package will have limited functionality. Please install it using the command: mim install 'mmcv>=2.0.1'"
-    )
-
-try:
-    from mmpose.apis import inference_topdown
-    from mmpose.apis import init_model as init_pose_estimator
-    from mmpose.evaluation.functional import nms
-    from mmpose.utils import adapt_mmdet_pipeline
-    from mmpose.structures import merge_data_samples
-except ImportError:
-    warnings.warn(
-        "The module 'mmpose' is not installed. The package will have limited functionality. Please install it using the command: mim install 'mmpose>=1.1.0'"
-    )
-        
-try:
-    from mmdet.apis import inference_detector, init_detector
-except ImportError:
-    warnings.warn(
-        "The module 'mmdet' is not installed. The package will have limited functionality. Please install it using the command: mim install 'mmdet>=3.1.0'"
-    )
-
-
-class Wholebody:
-    def __init__(self, 
-                 det_config=None, det_ckpt=None, 
-                 pose_config=None, pose_ckpt=None,
-                 device="cpu"):
-        
-        if det_config is None:
-            det_config = os.path.join(os.path.dirname(__file__), "yolox_config/yolox_l_8xb8-300e_coco.py")
-        
-        if pose_config is None:
-            pose_config = os.path.join(os.path.dirname(__file__), "dwpose_config/dwpose-l_384x288.py")
-
-        if det_ckpt is None:
-            det_ckpt = 'https://download.openmmlab.com/mmdetection/v2.0/yolox/yolox_l_8x8_300e_coco/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth'
-        
-        if pose_ckpt is None:
-            pose_ckpt = "https://huggingface.co/wanghaofan/dw-ll_ucoco_384/resolve/main/dw-ll_ucoco_384.pth"
-        
-        # build detector
-        self.detector = init_detector(det_config, det_ckpt, device=device)
-        self.detector.cfg = adapt_mmdet_pipeline(self.detector.cfg)
-
-        # build pose estimator
-        self.pose_estimator = init_pose_estimator(
-            pose_config,
-            pose_ckpt,
-            device=device)
-    
-    def to(self, device):
-        self.detector.to(device)
-        self.pose_estimator.to(device)
-        return self
-    
-    def __call__(self, oriImg):
-        # predict bbox
-        det_result = inference_detector(self.detector, oriImg)
-        pred_instance = det_result.pred_instances.cpu().numpy()
-        bboxes = np.concatenate(
-            (pred_instance.bboxes, pred_instance.scores[:, None]), axis=1)
-        bboxes = bboxes[np.logical_and(pred_instance.labels == 0,
-                                    pred_instance.scores > 0.5)]
-    
-        # set NMS threshold
-        bboxes = bboxes[nms(bboxes, 0.7), :4]
-
-        # predict keypoints
-        if len(bboxes) == 0:
-            pose_results = inference_topdown(self.pose_estimator, oriImg)
-        else:
-            pose_results = inference_topdown(self.pose_estimator, oriImg, bboxes)
-        preds = merge_data_samples(pose_results)
-        preds = preds.pred_instances
-
-        # preds = pose_results[0].pred_instances
-        keypoints = preds.get('transformed_keypoints',
-                                        preds.keypoints)
-        if 'keypoint_scores' in preds:
-            scores = preds.keypoint_scores
-        else:
-            scores = np.ones(keypoints.shape[:-1])
-
-        if 'keypoints_visible' in preds:
-            visible = preds.keypoints_visible
-        else:
-            visible = np.ones(keypoints.shape[:-1])
-        keypoints_info = np.concatenate(
-            (keypoints, scores[..., None], visible[..., None]),
-            axis=-1)
-        # compute neck joint
-        neck = np.mean(keypoints_info[:, [5, 6]], axis=1)
-        # neck score when visualizing pred
-        neck[:, 2:4] = np.logical_and(
-            keypoints_info[:, 5, 2:4] > 0.3,
-            keypoints_info[:, 6, 2:4] > 0.3).astype(int)
-        new_keypoints_info = np.insert(
-            keypoints_info, 17, neck, axis=1)
-        mmpose_idx = [
-            17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3
-        ]
-        openpose_idx = [
-            1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17
-        ]
-        new_keypoints_info[:, openpose_idx] = \
-            new_keypoints_info[:, mmpose_idx]
-        keypoints_info = new_keypoints_info
-
-        keypoints, scores, visible = keypoints_info[
-            ..., :2], keypoints_info[..., 2], keypoints_info[..., 3]
-        
-        return keypoints, scores

controlnet_aux_local/hed/__init__.py

@@ -1,129 +0,0 @@
-# This is an improved version and model of HED edge detection with Apache License, Version 2.0.
-# Please use this implementation in your products
-# This implementation may produce slightly different results from Saining Xie's official implementations,
-# but it generates smoother edges and is more suitable for ControlNet as well as other image-to-image translations.
-# Different from official models and other implementations, this is an RGB-input model (rather than BGR)
-# and in this way it works better for gradio's RGB protocol
-
-import os
-import warnings
-
-import cv2
-import numpy as np
-import torch
-from einops import rearrange
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from ..util import HWC3, nms, resize_image, safe_step
-
-
-class DoubleConvBlock(torch.nn.Module):
-    def __init__(self, input_channel, output_channel, layer_number):
-        super().__init__()
-        self.convs = torch.nn.Sequential()
-        self.convs.append(torch.nn.Conv2d(in_channels=input_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
-        for i in range(1, layer_number):
-            self.convs.append(torch.nn.Conv2d(in_channels=output_channel, out_channels=output_channel, kernel_size=(3, 3), stride=(1, 1), padding=1))
-        self.projection = torch.nn.Conv2d(in_channels=output_channel, out_channels=1, kernel_size=(1, 1), stride=(1, 1), padding=0)
-
-    def __call__(self, x, down_sampling=False):
-        h = x
-        if down_sampling:
-            h = torch.nn.functional.max_pool2d(h, kernel_size=(2, 2), stride=(2, 2))
-        for conv in self.convs:
-            h = conv(h)
-            h = torch.nn.functional.relu(h)
-        return h, self.projection(h)
-
-
-class ControlNetHED_Apache2(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.norm = torch.nn.Parameter(torch.zeros(size=(1, 3, 1, 1)))
-        self.block1 = DoubleConvBlock(input_channel=3, output_channel=64, layer_number=2)
-        self.block2 = DoubleConvBlock(input_channel=64, output_channel=128, layer_number=2)
-        self.block3 = DoubleConvBlock(input_channel=128, output_channel=256, layer_number=3)
-        self.block4 = DoubleConvBlock(input_channel=256, output_channel=512, layer_number=3)
-        self.block5 = DoubleConvBlock(input_channel=512, output_channel=512, layer_number=3)
-
-    def __call__(self, x):
-        h = x - self.norm
-        h, projection1 = self.block1(h)
-        h, projection2 = self.block2(h, down_sampling=True)
-        h, projection3 = self.block3(h, down_sampling=True)
-        h, projection4 = self.block4(h, down_sampling=True)
-        h, projection5 = self.block5(h, down_sampling=True)
-        return projection1, projection2, projection3, projection4, projection5
-
-class HEDdetector:
-    def __init__(self, netNetwork):
-        self.netNetwork = netNetwork
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None, local_files_only=False):
-        filename = filename or "ControlNetHED.pth"
-
-        if os.path.isdir(pretrained_model_or_path):
-            model_path = os.path.join(pretrained_model_or_path, filename)
-        else:
-            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir, local_files_only=local_files_only)
-
-        netNetwork = ControlNetHED_Apache2()
-        netNetwork.load_state_dict(torch.load(model_path, map_location='cpu'))
-        netNetwork.float().eval()
-
-        return cls(netNetwork)
-    
-    def to(self, device):
-        self.netNetwork.to(device)
-        return self
-    
-    def __call__(self, input_image, detect_resolution=512, image_resolution=512, safe=False, output_type="pil", scribble=False, **kwargs):
-        if "return_pil" in kwargs:
-            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
-            output_type = "pil" if kwargs["return_pil"] else "np"
-        if type(output_type) is bool:
-            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
-            if output_type:
-                output_type = "pil"
-
-        device = next(iter(self.netNetwork.parameters())).device
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-
-        assert input_image.ndim == 3
-        H, W, C = input_image.shape
-        with torch.no_grad():
-            image_hed = torch.from_numpy(input_image.copy()).float().to(device)
-            image_hed = rearrange(image_hed, 'h w c -> 1 c h w')
-            edges = self.netNetwork(image_hed)
-            edges = [e.detach().cpu().numpy().astype(np.float32)[0, 0] for e in edges]
-            edges = [cv2.resize(e, (W, H), interpolation=cv2.INTER_LINEAR) for e in edges]
-            edges = np.stack(edges, axis=2)
-            edge = 1 / (1 + np.exp(-np.mean(edges, axis=2).astype(np.float64)))
-            if safe:
-                edge = safe_step(edge)
-            edge = (edge * 255.0).clip(0, 255).astype(np.uint8)
-
-        detected_map = edge
-        detected_map = HWC3(detected_map)
-
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-        
-        if scribble:
-            detected_map = nms(detected_map, 127, 3.0)
-            detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
-            detected_map[detected_map > 4] = 255
-            detected_map[detected_map < 255] = 0
-
-        if output_type == "pil":
-            detected_map = Image.fromarray(detected_map)
-
-        return detected_map

controlnet_aux_local/leres/__init__.py

@@ -1,118 +0,0 @@
-import os
-
-import cv2
-import numpy as np
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from ..util import HWC3, resize_image
-from .leres.depthmap import estimateboost, estimateleres
-from .leres.multi_depth_model_woauxi import RelDepthModel
-from .leres.net_tools import strip_prefix_if_present
-from .pix2pix.models.pix2pix4depth_model import Pix2Pix4DepthModel
-from .pix2pix.options.test_options import TestOptions
-
-
-class LeresDetector:
-    def __init__(self, model, pix2pixmodel):
-        self.model = model
-        self.pix2pixmodel = pix2pixmodel
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_or_path, filename=None, pix2pix_filename=None, cache_dir=None, local_files_only=False):
-        filename = filename or "res101.pth"
-        pix2pix_filename = pix2pix_filename or "latest_net_G.pth"
-
-        if os.path.isdir(pretrained_model_or_path):
-            model_path = os.path.join(pretrained_model_or_path, filename)
-        else:
-            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir, local_files_only=local_files_only)
-            
-        checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
-
-        model = RelDepthModel(backbone='resnext101')
-        model.load_state_dict(strip_prefix_if_present(checkpoint['depth_model'], "module."), strict=True)
-        del checkpoint
-
-        if os.path.isdir(pretrained_model_or_path):
-            model_path = os.path.join(pretrained_model_or_path, pix2pix_filename)
-        else:
-            model_path = hf_hub_download(pretrained_model_or_path, pix2pix_filename, cache_dir=cache_dir, local_files_only=local_files_only)
-
-        opt = TestOptions().parse()
-        if not torch.cuda.is_available():
-            opt.gpu_ids = []  # cpu mode
-        pix2pixmodel = Pix2Pix4DepthModel(opt)
-        pix2pixmodel.save_dir = os.path.dirname(model_path)
-        pix2pixmodel.load_networks('latest')
-        pix2pixmodel.eval()
-
-        return cls(model, pix2pixmodel)
-
-    def to(self, device):
-        self.model.to(device)
-        # TODO - refactor pix2pix implementation to support device migration
-        # self.pix2pixmodel.to(device)
-        return self
-
-    def __call__(self, input_image, thr_a=0, thr_b=0, boost=False, detect_resolution=512, image_resolution=512, output_type="pil"):
-        device = next(iter(self.model.parameters())).device
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-        
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-
-        assert input_image.ndim == 3
-        height, width, dim = input_image.shape
-
-        with torch.no_grad():
-
-            if boost:
-                depth = estimateboost(input_image, self.model, 0, self.pix2pixmodel, max(width, height))
-            else:
-                depth = estimateleres(input_image, self.model, width, height)
-
-            numbytes=2
-            depth_min = depth.min()
-            depth_max = depth.max()
-            max_val = (2**(8*numbytes))-1
-
-            # check output before normalizing and mapping to 16 bit
-            if depth_max - depth_min > np.finfo("float").eps:
-                out = max_val * (depth - depth_min) / (depth_max - depth_min)
-            else:
-                out = np.zeros(depth.shape)
-            
-            # single channel, 16 bit image
-            depth_image = out.astype("uint16")
-
-            # convert to uint8
-            depth_image = cv2.convertScaleAbs(depth_image, alpha=(255.0/65535.0))
-
-            # remove near
-            if thr_a != 0:
-                thr_a = ((thr_a/100)*255) 
-                depth_image = cv2.threshold(depth_image, thr_a, 255, cv2.THRESH_TOZERO)[1]
-
-            # invert image
-            depth_image = cv2.bitwise_not(depth_image)
-
-            # remove bg
-            if thr_b != 0:
-                thr_b = ((thr_b/100)*255)
-                depth_image = cv2.threshold(depth_image, thr_b, 255, cv2.THRESH_TOZERO)[1]
-
-        detected_map = depth_image
-        detected_map = HWC3(detected_map)      
-
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-        
-        if output_type == "pil":
-            detected_map = Image.fromarray(detected_map)
-            
-        return detected_map

controlnet_aux_local/leres/leres/Resnet.py

@@ -1,199 +0,0 @@
-import torch.nn as nn
-import torch.nn as NN
-
-__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
-           'resnet152']
-
-
-model_urls = {
-    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
-    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
-    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
-    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
-    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
-}
-
-
-def conv3x3(in_planes, out_planes, stride=1):
-    """3x3 convolution with padding"""
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=1, bias=False)
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(BasicBlock, self).__init__()
-        self.conv1 = conv3x3(inplanes, planes, stride)
-        self.bn1 = NN.BatchNorm2d(planes) #NN.BatchNorm2d
-        self.relu = nn.ReLU(inplace=True)
-        self.conv2 = conv3x3(planes, planes)
-        self.bn2 = NN.BatchNorm2d(planes) #NN.BatchNorm2d
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(Bottleneck, self).__init__()
-        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
-        self.bn1 = NN.BatchNorm2d(planes) #NN.BatchNorm2d
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
-                               padding=1, bias=False)
-        self.bn2 = NN.BatchNorm2d(planes) #NN.BatchNorm2d
-        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
-        self.bn3 = NN.BatchNorm2d(planes * self.expansion) #NN.BatchNorm2d
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        out = self.bn3(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class ResNet(nn.Module):
-
-    def __init__(self, block, layers, num_classes=1000):
-        self.inplanes = 64
-        super(ResNet, self).__init__()
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
-                               bias=False)
-        self.bn1 = NN.BatchNorm2d(64)  #NN.BatchNorm2d
-        self.relu = nn.ReLU(inplace=True)
-        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-        self.layer1 = self._make_layer(block, 64, layers[0])
-        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
-        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
-        #self.avgpool = nn.AvgPool2d(7, stride=1)
-        #self.fc = nn.Linear(512 * block.expansion, num_classes)
-
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
-            elif isinstance(m, nn.BatchNorm2d):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
-
-    def _make_layer(self, block, planes, blocks, stride=1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=stride, bias=False),
-                NN.BatchNorm2d(planes * block.expansion), #NN.BatchNorm2d
-            )
-
-        layers = []
-        layers.append(block(self.inplanes, planes, stride, downsample))
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        features = []
-
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.relu(x)
-        x = self.maxpool(x)
-
-        x = self.layer1(x)
-        features.append(x)
-        x = self.layer2(x)
-        features.append(x)
-        x = self.layer3(x)
-        features.append(x)
-        x = self.layer4(x)
-        features.append(x)
-
-        return features
-
-
-def resnet18(pretrained=True, **kwargs):
-    """Constructs a ResNet-18 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
-    return model
-
-
-def resnet34(pretrained=True, **kwargs):
-    """Constructs a ResNet-34 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
-    return model
-
-
-def resnet50(pretrained=True, **kwargs):
-    """Constructs a ResNet-50 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
-
-    return model
-
-
-def resnet101(pretrained=True, **kwargs):
-    """Constructs a ResNet-101 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
-
-    return model
-
-
-def resnet152(pretrained=True, **kwargs):
-    """Constructs a ResNet-152 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
-    return model

controlnet_aux_local/leres/leres/Resnext_torch.py

@@ -1,237 +0,0 @@
-#!/usr/bin/env python
-# coding: utf-8
-import torch.nn as nn
-
-try:
-    from urllib import urlretrieve
-except ImportError:
-    from urllib.request import urlretrieve
-
-__all__ = ['resnext101_32x8d']
-
-
-model_urls = {
-    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
-    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
-}
-
-
-def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
-    """3x3 convolution with padding"""
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=dilation, groups=groups, bias=False, dilation=dilation)
-
-
-def conv1x1(in_planes, out_planes, stride=1):
-    """1x1 convolution"""
-    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
-                 base_width=64, dilation=1, norm_layer=None):
-        super(BasicBlock, self).__init__()
-        if norm_layer is None:
-            norm_layer = nn.BatchNorm2d
-        if groups != 1 or base_width != 64:
-            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
-        if dilation > 1:
-            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
-        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
-        self.conv1 = conv3x3(inplanes, planes, stride)
-        self.bn1 = norm_layer(planes)
-        self.relu = nn.ReLU(inplace=True)
-        self.conv2 = conv3x3(planes, planes)
-        self.bn2 = norm_layer(planes)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        identity = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        out += identity
-        out = self.relu(out)
-
-        return out
-
-
-class Bottleneck(nn.Module):
-    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
-    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
-    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
-    # This variant is also known as ResNet V1.5 and improves accuracy according to
-    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
-
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
-                 base_width=64, dilation=1, norm_layer=None):
-        super(Bottleneck, self).__init__()
-        if norm_layer is None:
-            norm_layer = nn.BatchNorm2d
-        width = int(planes * (base_width / 64.)) * groups
-        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
-        self.conv1 = conv1x1(inplanes, width)
-        self.bn1 = norm_layer(width)
-        self.conv2 = conv3x3(width, width, stride, groups, dilation)
-        self.bn2 = norm_layer(width)
-        self.conv3 = conv1x1(width, planes * self.expansion)
-        self.bn3 = norm_layer(planes * self.expansion)
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        identity = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        out = self.bn3(out)
-
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        out += identity
-        out = self.relu(out)
-
-        return out
-
-
-class ResNet(nn.Module):
-
-    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
-                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
-                 norm_layer=None):
-        super(ResNet, self).__init__()
-        if norm_layer is None:
-            norm_layer = nn.BatchNorm2d
-        self._norm_layer = norm_layer
-
-        self.inplanes = 64
-        self.dilation = 1
-        if replace_stride_with_dilation is None:
-            # each element in the tuple indicates if we should replace
-            # the 2x2 stride with a dilated convolution instead
-            replace_stride_with_dilation = [False, False, False]
-        if len(replace_stride_with_dilation) != 3:
-            raise ValueError("replace_stride_with_dilation should be None "
-                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
-        self.groups = groups
-        self.base_width = width_per_group
-        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
-                               bias=False)
-        self.bn1 = norm_layer(self.inplanes)
-        self.relu = nn.ReLU(inplace=True)
-        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-        self.layer1 = self._make_layer(block, 64, layers[0])
-        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
-                                       dilate=replace_stride_with_dilation[0])
-        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
-                                       dilate=replace_stride_with_dilation[1])
-        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
-                                       dilate=replace_stride_with_dilation[2])
-        #self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-        #self.fc = nn.Linear(512 * block.expansion, num_classes)
-
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
-            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
-
-        # Zero-initialize the last BN in each residual branch,
-        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
-        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
-        if zero_init_residual:
-            for m in self.modules():
-                if isinstance(m, Bottleneck):
-                    nn.init.constant_(m.bn3.weight, 0)
-                elif isinstance(m, BasicBlock):
-                    nn.init.constant_(m.bn2.weight, 0)
-
-    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
-        norm_layer = self._norm_layer
-        downsample = None
-        previous_dilation = self.dilation
-        if dilate:
-            self.dilation *= stride
-            stride = 1
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                conv1x1(self.inplanes, planes * block.expansion, stride),
-                norm_layer(planes * block.expansion),
-            )
-
-        layers = []
-        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
-                            self.base_width, previous_dilation, norm_layer))
-        self.inplanes = planes * block.expansion
-        for _ in range(1, blocks):
-            layers.append(block(self.inplanes, planes, groups=self.groups,
-                                base_width=self.base_width, dilation=self.dilation,
-                                norm_layer=norm_layer))
-
-        return nn.Sequential(*layers)
-
-    def _forward_impl(self, x):
-        # See note [TorchScript super()]
-        features = []
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.relu(x)
-        x = self.maxpool(x)
-
-        x = self.layer1(x)
-        features.append(x)
-
-        x = self.layer2(x)
-        features.append(x)
-
-        x = self.layer3(x)
-        features.append(x)
-
-        x = self.layer4(x)
-        features.append(x)
-
-        #x = self.avgpool(x)
-        #x = torch.flatten(x, 1)
-        #x = self.fc(x)
-
-        return features
-
-    def forward(self, x):
-        return self._forward_impl(x)
-
-
-
-def resnext101_32x8d(pretrained=True, **kwargs):
-    """Constructs a ResNet-152 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    kwargs['groups'] = 32
-    kwargs['width_per_group'] = 8
-
-    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
-    return model
-

controlnet_aux_local/leres/leres/depthmap.py

@@ -1,548 +0,0 @@
-# Author: thygate
-# https://github.com/thygate/stable-diffusion-webui-depthmap-script
-
-import gc
-from operator import getitem
-
-import cv2
-import numpy as np
-import skimage.measure
-import torch
-from torchvision.transforms import transforms
-
-from ...util import torch_gc
-
-whole_size_threshold = 1600  # R_max from the paper
-pix2pixsize = 1024
-
-def scale_torch(img):
-    """
-    Scale the image and output it in torch.tensor.
-    :param img: input rgb is in shape [H, W, C], input depth/disp is in shape [H, W]
-    :param scale: the scale factor. float
-    :return: img. [C, H, W]
-    """
-    if len(img.shape) == 2:
-        img = img[np.newaxis, :, :]
-    if img.shape[2] == 3:
-        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406) , (0.229, 0.224, 0.225) )])
-        img = transform(img.astype(np.float32))
-    else:
-        img = img.astype(np.float32)
-        img = torch.from_numpy(img)
-    return img
-    
-def estimateleres(img, model, w, h):
-    device = next(iter(model.parameters())).device
-    # leres transform input
-    rgb_c = img[:, :, ::-1].copy()
-    A_resize = cv2.resize(rgb_c, (w, h))
-    img_torch = scale_torch(A_resize)[None, :, :, :] 
-    
-    # compute
-    with torch.no_grad():
-        img_torch = img_torch.to(device)
-        prediction = model.depth_model(img_torch)
-
-    prediction = prediction.squeeze().cpu().numpy()
-    prediction = cv2.resize(prediction, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)
-
-    return prediction
-
-def generatemask(size):
-    # Generates a Guassian mask
-    mask = np.zeros(size, dtype=np.float32)
-    sigma = int(size[0]/16)
-    k_size = int(2 * np.ceil(2 * int(size[0]/16)) + 1)
-    mask[int(0.15*size[0]):size[0] - int(0.15*size[0]), int(0.15*size[1]): size[1] - int(0.15*size[1])] = 1
-    mask = cv2.GaussianBlur(mask, (int(k_size), int(k_size)), sigma)
-    mask = (mask - mask.min()) / (mask.max() - mask.min())
-    mask = mask.astype(np.float32)
-    return mask
-
-def resizewithpool(img, size):
-    i_size = img.shape[0]
-    n = int(np.floor(i_size/size))
-
-    out = skimage.measure.block_reduce(img, (n, n), np.max)
-    return out
-
-def rgb2gray(rgb):
-    # Converts rgb to gray
-    return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140])
-
-def calculateprocessingres(img, basesize, confidence=0.1, scale_threshold=3, whole_size_threshold=3000):
-    # Returns the R_x resolution described in section 5 of the main paper.
-
-    # Parameters:
-    #    img :input rgb image
-    #    basesize : size the dilation kernel which is equal to receptive field of the network.
-    #    confidence: value of x in R_x; allowed percentage of pixels that are not getting any contextual cue.
-    #    scale_threshold: maximum allowed upscaling on the input image ; it has been set to 3.
-    #    whole_size_threshold: maximum allowed resolution. (R_max from section 6 of the main paper)
-
-    # Returns:
-    #    outputsize_scale*speed_scale :The computed R_x resolution
-    #    patch_scale: K parameter from section 6 of the paper
-
-    # speed scale parameter is to process every image in a smaller size to accelerate the R_x resolution search
-    speed_scale = 32
-    image_dim = int(min(img.shape[0:2]))
-
-    gray = rgb2gray(img)
-    grad = np.abs(cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)) + np.abs(cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3))
-    grad = cv2.resize(grad, (image_dim, image_dim), cv2.INTER_AREA)
-
-    # thresholding the gradient map to generate the edge-map as a proxy of the contextual cues
-    m = grad.min()
-    M = grad.max()
-    middle = m + (0.4 * (M - m))
-    grad[grad < middle] = 0
-    grad[grad >= middle] = 1
-
-    # dilation kernel with size of the receptive field
-    kernel = np.ones((int(basesize/speed_scale), int(basesize/speed_scale)), float)
-    # dilation kernel with size of the a quarter of receptive field used to compute k
-    # as described in section 6 of main paper
-    kernel2 = np.ones((int(basesize / (4*speed_scale)), int(basesize / (4*speed_scale))), float)
-
-    # Output resolution limit set by the whole_size_threshold and scale_threshold.
-    threshold = min(whole_size_threshold, scale_threshold * max(img.shape[:2]))
-
-    outputsize_scale = basesize / speed_scale
-    for p_size in range(int(basesize/speed_scale), int(threshold/speed_scale), int(basesize / (2*speed_scale))):
-        grad_resized = resizewithpool(grad, p_size)
-        grad_resized = cv2.resize(grad_resized, (p_size, p_size), cv2.INTER_NEAREST)
-        grad_resized[grad_resized >= 0.5] = 1
-        grad_resized[grad_resized < 0.5] = 0
-
-        dilated = cv2.dilate(grad_resized, kernel, iterations=1)
-        meanvalue = (1-dilated).mean()
-        if meanvalue > confidence:
-            break
-        else:
-            outputsize_scale = p_size
-
-    grad_region = cv2.dilate(grad_resized, kernel2, iterations=1)
-    patch_scale = grad_region.mean()
-
-    return int(outputsize_scale*speed_scale), patch_scale
-
-# Generate a double-input depth estimation
-def doubleestimate(img, size1, size2, pix2pixsize, model, net_type, pix2pixmodel):
-    # Generate the low resolution estimation
-    estimate1 = singleestimate(img, size1, model, net_type)
-    # Resize to the inference size of merge network.
-    estimate1 = cv2.resize(estimate1, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
-
-    # Generate the high resolution estimation
-    estimate2 = singleestimate(img, size2, model, net_type)
-    # Resize to the inference size of merge network.
-    estimate2 = cv2.resize(estimate2, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
-
-    # Inference on the merge model
-    pix2pixmodel.set_input(estimate1, estimate2)
-    pix2pixmodel.test()
-    visuals = pix2pixmodel.get_current_visuals()
-    prediction_mapped = visuals['fake_B']
-    prediction_mapped = (prediction_mapped+1)/2
-    prediction_mapped = (prediction_mapped - torch.min(prediction_mapped)) / (
-                torch.max(prediction_mapped) - torch.min(prediction_mapped))
-    prediction_mapped = prediction_mapped.squeeze().cpu().numpy()
-
-    return prediction_mapped
-
-# Generate a single-input depth estimation
-def singleestimate(img, msize, model, net_type):
-    # if net_type == 0:
-    return estimateleres(img, model, msize, msize)
-    # else:
-    # 	return estimatemidasBoost(img, model, msize, msize)
-
-def applyGridpatch(blsize, stride, img, box):
-    # Extract a simple grid patch.
-    counter1 = 0
-    patch_bound_list = {}
-    for k in range(blsize, img.shape[1] - blsize, stride):
-        for j in range(blsize, img.shape[0] - blsize, stride):
-            patch_bound_list[str(counter1)] = {}
-            patchbounds = [j - blsize, k - blsize, j - blsize + 2 * blsize, k - blsize + 2 * blsize]
-            patch_bound = [box[0] + patchbounds[1], box[1] + patchbounds[0], patchbounds[3] - patchbounds[1],
-                           patchbounds[2] - patchbounds[0]]
-            patch_bound_list[str(counter1)]['rect'] = patch_bound
-            patch_bound_list[str(counter1)]['size'] = patch_bound[2]
-            counter1 = counter1 + 1
-    return patch_bound_list
-
-# Generating local patches to perform the local refinement described in section 6 of the main paper.
-def generatepatchs(img, base_size):
-    
-    # Compute the gradients as a proxy of the contextual cues.
-    img_gray = rgb2gray(img)
-    whole_grad = np.abs(cv2.Sobel(img_gray, cv2.CV_64F, 0, 1, ksize=3)) +\
-        np.abs(cv2.Sobel(img_gray, cv2.CV_64F, 1, 0, ksize=3))
-
-    threshold = whole_grad[whole_grad > 0].mean()
-    whole_grad[whole_grad < threshold] = 0
-
-    # We use the integral image to speed-up the evaluation of the amount of gradients for each patch.
-    gf = whole_grad.sum()/len(whole_grad.reshape(-1))
-    grad_integral_image = cv2.integral(whole_grad)
-
-    # Variables are selected such that the initial patch size would be the receptive field size
-    # and the stride is set to 1/3 of the receptive field size.
-    blsize = int(round(base_size/2))
-    stride = int(round(blsize*0.75))
-
-    # Get initial Grid
-    patch_bound_list = applyGridpatch(blsize, stride, img, [0, 0, 0, 0])
-
-    # Refine initial Grid of patches by discarding the flat (in terms of gradients of the rgb image) ones. Refine
-    # each patch size to ensure that there will be enough depth cues for the network to generate a consistent depth map.
-    print("Selecting patches ...")
-    patch_bound_list = adaptiveselection(grad_integral_image, patch_bound_list, gf)
-
-    # Sort the patch list to make sure the merging operation will be done with the correct order: starting from biggest
-    # patch
-    patchset = sorted(patch_bound_list.items(), key=lambda x: getitem(x[1], 'size'), reverse=True)
-    return patchset
-
-def getGF_fromintegral(integralimage, rect):
-    # Computes the gradient density of a given patch from the gradient integral image.
-    x1 = rect[1]
-    x2 = rect[1]+rect[3]
-    y1 = rect[0]
-    y2 = rect[0]+rect[2]
-    value = integralimage[x2, y2]-integralimage[x1, y2]-integralimage[x2, y1]+integralimage[x1, y1]
-    return value
-
-# Adaptively select patches
-def adaptiveselection(integral_grad, patch_bound_list, gf):
-    patchlist = {}
-    count = 0
-    height, width = integral_grad.shape
-
-    search_step = int(32/factor)
-
-    # Go through all patches
-    for c in range(len(patch_bound_list)):
-        # Get patch
-        bbox = patch_bound_list[str(c)]['rect']
-
-        # Compute the amount of gradients present in the patch from the integral image.
-        cgf = getGF_fromintegral(integral_grad, bbox)/(bbox[2]*bbox[3])
-
-        # Check if patching is beneficial by comparing the gradient density of the patch to
-        # the gradient density of the whole image
-        if cgf >= gf:
-            bbox_test = bbox.copy()
-            patchlist[str(count)] = {}
-
-            # Enlarge each patch until the gradient density of the patch is equal
-            # to the whole image gradient density
-            while True:
-
-                bbox_test[0] = bbox_test[0] - int(search_step/2)
-                bbox_test[1] = bbox_test[1] - int(search_step/2)
-
-                bbox_test[2] = bbox_test[2] + search_step
-                bbox_test[3] = bbox_test[3] + search_step
-
-                # Check if we are still within the image
-                if bbox_test[0] < 0 or bbox_test[1] < 0 or bbox_test[1] + bbox_test[3] >= height \
-                        or bbox_test[0] + bbox_test[2] >= width:
-                    break
-
-                # Compare gradient density
-                cgf = getGF_fromintegral(integral_grad, bbox_test)/(bbox_test[2]*bbox_test[3])
-                if cgf < gf:
-                    break
-                bbox = bbox_test.copy()
-
-            # Add patch to selected patches
-            patchlist[str(count)]['rect'] = bbox
-            patchlist[str(count)]['size'] = bbox[2]
-            count = count + 1
-    
-    # Return selected patches
-    return patchlist
-
-def impatch(image, rect):
-    # Extract the given patch pixels from a given image.
-    w1 = rect[0]
-    h1 = rect[1]
-    w2 = w1 + rect[2]
-    h2 = h1 + rect[3]
-    image_patch = image[h1:h2, w1:w2]
-    return image_patch
-
-class ImageandPatchs:
-    def __init__(self, root_dir, name, patchsinfo, rgb_image, scale=1):
-        self.root_dir = root_dir
-        self.patchsinfo = patchsinfo
-        self.name = name
-        self.patchs = patchsinfo
-        self.scale = scale
-
-        self.rgb_image = cv2.resize(rgb_image, (round(rgb_image.shape[1]*scale), round(rgb_image.shape[0]*scale)),
-                                    interpolation=cv2.INTER_CUBIC)
-
-        self.do_have_estimate = False
-        self.estimation_updated_image = None
-        self.estimation_base_image = None
-
-    def __len__(self):
-        return len(self.patchs)
-
-    def set_base_estimate(self, est):
-        self.estimation_base_image = est
-        if self.estimation_updated_image is not None:
-            self.do_have_estimate = True
-
-    def set_updated_estimate(self, est):
-        self.estimation_updated_image = est
-        if self.estimation_base_image is not None:
-            self.do_have_estimate = True
-
-    def __getitem__(self, index):
-        patch_id = int(self.patchs[index][0])
-        rect = np.array(self.patchs[index][1]['rect'])
-        msize = self.patchs[index][1]['size']
-
-        ## applying scale to rect:
-        rect = np.round(rect * self.scale)
-        rect = rect.astype('int')
-        msize = round(msize * self.scale)
-
-        patch_rgb = impatch(self.rgb_image, rect)
-        if self.do_have_estimate:
-            patch_whole_estimate_base = impatch(self.estimation_base_image, rect)
-            patch_whole_estimate_updated = impatch(self.estimation_updated_image, rect)
-            return {'patch_rgb': patch_rgb, 'patch_whole_estimate_base': patch_whole_estimate_base,
-                    'patch_whole_estimate_updated': patch_whole_estimate_updated, 'rect': rect,
-                    'size': msize, 'id': patch_id}
-        else:
-            return {'patch_rgb': patch_rgb, 'rect': rect, 'size': msize, 'id': patch_id}
-
-    def print_options(self, opt):
-        """Print and save options
-
-        It will print both current options and default values(if different).
-        It will save options into a text file / [checkpoints_dir] / opt.txt
-        """
-        message = ''
-        message += '----------------- Options ---------------\n'
-        for k, v in sorted(vars(opt).items()):
-            comment = ''
-            default = self.parser.get_default(k)
-            if v != default:
-                comment = '\t[default: %s]' % str(default)
-            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
-        message += '----------------- End -------------------'
-        print(message)
-
-        # save to the disk
-        """
-        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
-        util.mkdirs(expr_dir)
-        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
-        with open(file_name, 'wt') as opt_file:
-            opt_file.write(message)
-            opt_file.write('\n')
-        """
-
-    def parse(self):
-        """Parse our options, create checkpoints directory suffix, and set up gpu device."""
-        opt = self.gather_options()
-        opt.isTrain = self.isTrain   # train or test
-
-        # process opt.suffix
-        if opt.suffix:
-            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
-            opt.name = opt.name + suffix
-
-        #self.print_options(opt)
-
-        # set gpu ids
-        str_ids = opt.gpu_ids.split(',')
-        opt.gpu_ids = []
-        for str_id in str_ids:
-            id = int(str_id)
-            if id >= 0:
-                opt.gpu_ids.append(id)
-        #if len(opt.gpu_ids) > 0:
-        #    torch.cuda.set_device(opt.gpu_ids[0])
-
-        self.opt = opt
-        return self.opt
-
-
-def estimateboost(img, model, model_type, pix2pixmodel, max_res=512, depthmap_script_boost_rmax=None):
-    global whole_size_threshold
-    
-    # get settings
-    if depthmap_script_boost_rmax:
-        whole_size_threshold = depthmap_script_boost_rmax
-        
-    if model_type == 0: #leres
-        net_receptive_field_size = 448
-        patch_netsize = 2 * net_receptive_field_size
-    elif model_type == 1: #dpt_beit_large_512
-        net_receptive_field_size = 512
-        patch_netsize = 2 * net_receptive_field_size
-    else: #other midas
-        net_receptive_field_size = 384
-        patch_netsize = 2 * net_receptive_field_size
-
-    gc.collect()
-    torch_gc()
-
-    # Generate mask used to smoothly blend the local pathc estimations to the base estimate.
-    # It is arbitrarily large to avoid artifacts during rescaling for each crop.
-    mask_org = generatemask((3000, 3000))
-    mask = mask_org.copy()
-
-    # Value x of R_x defined in the section 5 of the main paper.
-    r_threshold_value = 0.2
-    #if R0:
-    #	r_threshold_value = 0
-
-    input_resolution = img.shape
-    scale_threshold = 3  # Allows up-scaling with a scale up to 3
-
-    # Find the best input resolution R-x. The resolution search described in section 5-double estimation of the main paper and section B of the
-    # supplementary material.
-    whole_image_optimal_size, patch_scale = calculateprocessingres(img, net_receptive_field_size, r_threshold_value, scale_threshold, whole_size_threshold)
-
-    # print('wholeImage being processed in :', whole_image_optimal_size)
-
-    # Generate the base estimate using the double estimation.
-    whole_estimate = doubleestimate(img, net_receptive_field_size, whole_image_optimal_size, pix2pixsize, model, model_type, pix2pixmodel)
-
-    # Compute the multiplier described in section 6 of the main paper to make sure our initial patch can select
-    # small high-density regions of the image.
-    global factor
-    factor = max(min(1, 4 * patch_scale * whole_image_optimal_size / whole_size_threshold), 0.2)
-    # print('Adjust factor is:', 1/factor)
-        
-    # Check if Local boosting is beneficial.
-    if max_res < whole_image_optimal_size:
-        # print("No Local boosting. Specified Max Res is smaller than R20, Returning doubleestimate result")
-        return cv2.resize(whole_estimate, (input_resolution[1], input_resolution[0]), interpolation=cv2.INTER_CUBIC)
-
-    # Compute the default target resolution.
-    if img.shape[0] > img.shape[1]:
-        a = 2 * whole_image_optimal_size
-        b = round(2 * whole_image_optimal_size * img.shape[1] / img.shape[0])
-    else:
-        a = round(2 * whole_image_optimal_size * img.shape[0] / img.shape[1])
-        b = 2 * whole_image_optimal_size
-    b = int(round(b / factor))
-    a = int(round(a / factor))
-
-    """
-    # recompute a, b and saturate to max res.
-    if max(a,b) > max_res:
-        print('Default Res is higher than max-res: Reducing final resolution')
-        if img.shape[0] > img.shape[1]:
-            a = max_res
-            b = round(max_res * img.shape[1] / img.shape[0])
-        else:
-            a = round(max_res * img.shape[0] / img.shape[1])
-            b = max_res
-        b = int(b)
-        a = int(a)
-    """
-
-    img = cv2.resize(img, (b, a), interpolation=cv2.INTER_CUBIC)
-
-    # Extract selected patches for local refinement
-    base_size = net_receptive_field_size * 2
-    patchset = generatepatchs(img, base_size)
-
-    # print('Target resolution: ', img.shape)
-
-    # Computing a scale in case user prompted to generate the results as the same resolution of the input.
-    # Notice that our method output resolution is independent of the input resolution and this parameter will only
-    # enable a scaling operation during the local patch merge implementation to generate results with the same resolution
-    # as the input.
-    """
-    if output_resolution == 1:
-        mergein_scale = input_resolution[0] / img.shape[0]
-        print('Dynamicly change merged-in resolution; scale:', mergein_scale)
-    else:
-        mergein_scale = 1
-    """
-    # always rescale to input res for now
-    mergein_scale = input_resolution[0] / img.shape[0]
-
-    imageandpatchs = ImageandPatchs('', '', patchset, img, mergein_scale)
-    whole_estimate_resized = cv2.resize(whole_estimate, (round(img.shape[1]*mergein_scale),
-                                        round(img.shape[0]*mergein_scale)), interpolation=cv2.INTER_CUBIC)
-    imageandpatchs.set_base_estimate(whole_estimate_resized.copy())
-    imageandpatchs.set_updated_estimate(whole_estimate_resized.copy())
-
-    print('Resulting depthmap resolution will be :', whole_estimate_resized.shape[:2])
-    print('Patches to process: '+str(len(imageandpatchs)))
-
-    # Enumerate through all patches, generate their estimations and refining the base estimate.
-    for patch_ind in range(len(imageandpatchs)):
-        
-        # Get patch information
-        patch = imageandpatchs[patch_ind] # patch object
-        patch_rgb = patch['patch_rgb'] # rgb patch
-        patch_whole_estimate_base = patch['patch_whole_estimate_base'] # corresponding patch from base
-        rect = patch['rect'] # patch size and location
-        patch_id = patch['id'] # patch ID
-        org_size = patch_whole_estimate_base.shape # the original size from the unscaled input
-        print('\t Processing patch', patch_ind, '/', len(imageandpatchs)-1, '|', rect)
-
-        # We apply double estimation for patches. The high resolution value is fixed to twice the receptive
-        # field size of the network for patches to accelerate the process.
-        patch_estimation = doubleestimate(patch_rgb, net_receptive_field_size, patch_netsize, pix2pixsize, model, model_type, pix2pixmodel)
-        patch_estimation = cv2.resize(patch_estimation, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
-        patch_whole_estimate_base = cv2.resize(patch_whole_estimate_base, (pix2pixsize, pix2pixsize), interpolation=cv2.INTER_CUBIC)
-
-        # Merging the patch estimation into the base estimate using our merge network:
-        # We feed the patch estimation and the same region from the updated base estimate to the merge network
-        # to generate the target estimate for the corresponding region.
-        pix2pixmodel.set_input(patch_whole_estimate_base, patch_estimation)
-
-        # Run merging network
-        pix2pixmodel.test()
-        visuals = pix2pixmodel.get_current_visuals()
-
-        prediction_mapped = visuals['fake_B']
-        prediction_mapped = (prediction_mapped+1)/2
-        prediction_mapped = prediction_mapped.squeeze().cpu().numpy()
-
-        mapped = prediction_mapped
-
-        # We use a simple linear polynomial to make sure the result of the merge network would match the values of
-        # base estimate
-        p_coef = np.polyfit(mapped.reshape(-1), patch_whole_estimate_base.reshape(-1), deg=1)
-        merged = np.polyval(p_coef, mapped.reshape(-1)).reshape(mapped.shape)
-
-        merged = cv2.resize(merged, (org_size[1],org_size[0]), interpolation=cv2.INTER_CUBIC)
-
-        # Get patch size and location
-        w1 = rect[0]
-        h1 = rect[1]
-        w2 = w1 + rect[2]
-        h2 = h1 + rect[3]
-
-        # To speed up the implementation, we only generate the Gaussian mask once with a sufficiently large size
-        # and resize it to our needed size while merging the patches.
-        if mask.shape != org_size:
-            mask = cv2.resize(mask_org, (org_size[1],org_size[0]), interpolation=cv2.INTER_LINEAR)
-
-        tobemergedto = imageandpatchs.estimation_updated_image
-
-        # Update the whole estimation:
-        # We use a simple Gaussian mask to blend the merged patch region with the base estimate to ensure seamless
-        # blending at the boundaries of the patch region.
-        tobemergedto[h1:h2, w1:w2] = np.multiply(tobemergedto[h1:h2, w1:w2], 1 - mask) + np.multiply(merged, mask)
-        imageandpatchs.set_updated_estimate(tobemergedto)
-
-    # output
-    return cv2.resize(imageandpatchs.estimation_updated_image, (input_resolution[1], input_resolution[0]), interpolation=cv2.INTER_CUBIC)

controlnet_aux_local/leres/leres/multi_depth_model_woauxi.py

@@ -1,35 +0,0 @@
-import torch
-import torch.nn as nn
-
-from . import network_auxi as network
-from .net_tools import get_func
-
-
-class RelDepthModel(nn.Module):
-    def __init__(self, backbone='resnet50'):
-        super(RelDepthModel, self).__init__()
-        if backbone == 'resnet50':
-            encoder = 'resnet50_stride32'
-        elif backbone == 'resnext101':
-            encoder = 'resnext101_stride32x8d'
-        self.depth_model = DepthModel(encoder)
-
-    def inference(self, rgb):
-        with torch.no_grad():
-            input = rgb.to(self.depth_model.device)
-            depth = self.depth_model(input)
-            #pred_depth_out = depth - depth.min() + 0.01
-            return depth #pred_depth_out
-
-
-class DepthModel(nn.Module):
-    def __init__(self, encoder):
-        super(DepthModel, self).__init__()
-        backbone = network.__name__.split('.')[-1] + '.' + encoder
-        self.encoder_modules = get_func(backbone)()
-        self.decoder_modules = network.Decoder()
-
-    def forward(self, x):
-        lateral_out = self.encoder_modules(x)
-        out_logit = self.decoder_modules(lateral_out)
-        return out_logit

controlnet_aux_local/leres/leres/net_tools.py

@@ -1,54 +0,0 @@
-import importlib
-import torch
-import os
-from collections import OrderedDict
-
-
-def get_func(func_name):
-    """Helper to return a function object by name. func_name must identify a
-    function in this module or the path to a function relative to the base
-    'modeling' module.
-    """
-    if func_name == '':
-        return None
-    try:
-        parts = func_name.split('.')
-        # Refers to a function in this module
-        if len(parts) == 1:
-            return globals()[parts[0]]
-        # Otherwise, assume we're referencing a module under modeling
-        module_name = 'controlnet_aux.leres.leres.' + '.'.join(parts[:-1])
-        module = importlib.import_module(module_name)
-        return getattr(module, parts[-1])
-    except Exception:
-        print('Failed to f1ind function: %s', func_name)
-        raise
-
-def load_ckpt(args, depth_model, shift_model, focal_model):
-    """
-    Load checkpoint.
-    """
-    if os.path.isfile(args.load_ckpt):
-        print("loading checkpoint %s" % args.load_ckpt)
-        checkpoint = torch.load(args.load_ckpt)
-        if shift_model is not None:
-            shift_model.load_state_dict(strip_prefix_if_present(checkpoint['shift_model'], 'module.'),
-                                    strict=True)
-        if focal_model is not None:
-            focal_model.load_state_dict(strip_prefix_if_present(checkpoint['focal_model'], 'module.'),
-                                    strict=True)
-        depth_model.load_state_dict(strip_prefix_if_present(checkpoint['depth_model'], "module."),
-                                    strict=True)
-        del checkpoint
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-
-def strip_prefix_if_present(state_dict, prefix):
-    keys = sorted(state_dict.keys())
-    if not all(key.startswith(prefix) for key in keys):
-        return state_dict
-    stripped_state_dict = OrderedDict()
-    for key, value in state_dict.items():
-        stripped_state_dict[key.replace(prefix, "")] = value
-    return stripped_state_dict

controlnet_aux_local/leres/leres/network_auxi.py

@@ -1,417 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.init as init
-
-from . import Resnet, Resnext_torch
-
-
-def resnet50_stride32():
-    return DepthNet(backbone='resnet', depth=50, upfactors=[2, 2, 2, 2])
-
-def resnext101_stride32x8d():
-    return DepthNet(backbone='resnext101_32x8d', depth=101, upfactors=[2, 2, 2, 2])
-
-
-class Decoder(nn.Module):
-    def __init__(self):
-        super(Decoder, self).__init__()
-        self.inchannels =  [256, 512, 1024, 2048]
-        self.midchannels = [256, 256, 256, 512]
-        self.upfactors = [2,2,2,2]
-        self.outchannels = 1
-
-        self.conv = FTB(inchannels=self.inchannels[3], midchannels=self.midchannels[3])
-        self.conv1 = nn.Conv2d(in_channels=self.midchannels[3], out_channels=self.midchannels[2], kernel_size=3, padding=1, stride=1, bias=True)
-        self.upsample = nn.Upsample(scale_factor=self.upfactors[3], mode='bilinear', align_corners=True)
-        
-        self.ffm2 = FFM(inchannels=self.inchannels[2], midchannels=self.midchannels[2], outchannels = self.midchannels[2], upfactor=self.upfactors[2])
-        self.ffm1 = FFM(inchannels=self.inchannels[1], midchannels=self.midchannels[1], outchannels = self.midchannels[1], upfactor=self.upfactors[1])
-        self.ffm0 = FFM(inchannels=self.inchannels[0], midchannels=self.midchannels[0], outchannels = self.midchannels[0], upfactor=self.upfactors[0])
-        
-        self.outconv = AO(inchannels=self.midchannels[0], outchannels=self.outchannels, upfactor=2)
-        self._init_params()
-        
-    def _init_params(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.ConvTranspose2d):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm2d): #NN.BatchNorm2d
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-                    
-    def forward(self, features):
-        x_32x = self.conv(features[3])  # 1/32
-        x_32 = self.conv1(x_32x)
-        x_16 = self.upsample(x_32)  # 1/16
-
-        x_8 = self.ffm2(features[2], x_16)  # 1/8
-        x_4 = self.ffm1(features[1], x_8)  # 1/4
-        x_2 = self.ffm0(features[0], x_4)  # 1/2
-        #-----------------------------------------
-        x = self.outconv(x_2)  # original size
-        return x
-
-class DepthNet(nn.Module):
-    __factory = {
-        18: Resnet.resnet18,
-        34: Resnet.resnet34,
-        50: Resnet.resnet50,
-        101: Resnet.resnet101,
-        152: Resnet.resnet152
-    }
-    def __init__(self,
-                 backbone='resnet',
-                 depth=50,
-                 upfactors=[2, 2, 2, 2]):
-        super(DepthNet, self).__init__()
-        self.backbone = backbone
-        self.depth = depth
-        self.pretrained = False
-        self.inchannels = [256, 512, 1024, 2048]
-        self.midchannels = [256, 256, 256, 512]
-        self.upfactors = upfactors
-        self.outchannels = 1
-
-        # Build model
-        if self.backbone == 'resnet':
-            if self.depth not in DepthNet.__factory:
-                raise KeyError("Unsupported depth:", self.depth)
-            self.encoder = DepthNet.__factory[depth](pretrained=self.pretrained)
-        elif self.backbone == 'resnext101_32x8d':
-            self.encoder = Resnext_torch.resnext101_32x8d(pretrained=self.pretrained)
-        else:
-            self.encoder = Resnext_torch.resnext101(pretrained=self.pretrained)
-
-    def forward(self, x):
-        x = self.encoder(x)  # 1/32, 1/16, 1/8, 1/4
-        return x
-
-
-class FTB(nn.Module):
-    def __init__(self, inchannels, midchannels=512):
-        super(FTB, self).__init__()
-        self.in1 = inchannels
-        self.mid = midchannels
-        self.conv1 = nn.Conv2d(in_channels=self.in1, out_channels=self.mid, kernel_size=3, padding=1, stride=1,
-                               bias=True)
-        # NN.BatchNorm2d
-        self.conv_branch = nn.Sequential(nn.ReLU(inplace=True), \
-                                         nn.Conv2d(in_channels=self.mid, out_channels=self.mid, kernel_size=3,
-                                                   padding=1, stride=1, bias=True), \
-                                         nn.BatchNorm2d(num_features=self.mid), \
-                                         nn.ReLU(inplace=True), \
-                                         nn.Conv2d(in_channels=self.mid, out_channels=self.mid, kernel_size=3,
-                                                   padding=1, stride=1, bias=True))
-        self.relu = nn.ReLU(inplace=True)
-
-        self.init_params()
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = x + self.conv_branch(x)
-        x = self.relu(x)
-
-        return x
-
-    def init_params(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.ConvTranspose2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm2d):  # NN.BatchNorm2d
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-
-
-class ATA(nn.Module):
-    def __init__(self, inchannels, reduction=8):
-        super(ATA, self).__init__()
-        self.inchannels = inchannels
-        self.avg_pool = nn.AdaptiveAvgPool2d(1)
-        self.fc = nn.Sequential(nn.Linear(self.inchannels * 2, self.inchannels // reduction),
-                                nn.ReLU(inplace=True),
-                                nn.Linear(self.inchannels // reduction, self.inchannels),
-                                nn.Sigmoid())
-        self.init_params()
-
-    def forward(self, low_x, high_x):
-        n, c, _, _ = low_x.size()
-        x = torch.cat([low_x, high_x], 1)
-        x = self.avg_pool(x)
-        x = x.view(n, -1)
-        x = self.fc(x).view(n, c, 1, 1)
-        x = low_x * x + high_x
-
-        return x
-
-    def init_params(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                # init.normal(m.weight, std=0.01)
-                init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.ConvTranspose2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                # init.normal_(m.weight, std=0.01)
-                init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm2d):  # NN.BatchNorm2d
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-
-
-class FFM(nn.Module):
-    def __init__(self, inchannels, midchannels, outchannels, upfactor=2):
-        super(FFM, self).__init__()
-        self.inchannels = inchannels
-        self.midchannels = midchannels
-        self.outchannels = outchannels
-        self.upfactor = upfactor
-
-        self.ftb1 = FTB(inchannels=self.inchannels, midchannels=self.midchannels)
-        # self.ata = ATA(inchannels = self.midchannels)
-        self.ftb2 = FTB(inchannels=self.midchannels, midchannels=self.outchannels)
-
-        self.upsample = nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True)
-
-        self.init_params()
-
-    def forward(self, low_x, high_x):
-        x = self.ftb1(low_x)
-        x = x + high_x
-        x = self.ftb2(x)
-        x = self.upsample(x)
-
-        return x
-
-    def init_params(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.ConvTranspose2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm2d):  # NN.Batchnorm2d
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-
-
-class AO(nn.Module):
-    # Adaptive output module
-    def __init__(self, inchannels, outchannels, upfactor=2):
-        super(AO, self).__init__()
-        self.inchannels = inchannels
-        self.outchannels = outchannels
-        self.upfactor = upfactor
-
-        self.adapt_conv = nn.Sequential(
-            nn.Conv2d(in_channels=self.inchannels, out_channels=self.inchannels // 2, kernel_size=3, padding=1,
-                      stride=1, bias=True), \
-            nn.BatchNorm2d(num_features=self.inchannels // 2), \
-            nn.ReLU(inplace=True), \
-            nn.Conv2d(in_channels=self.inchannels // 2, out_channels=self.outchannels, kernel_size=3, padding=1,
-                      stride=1, bias=True), \
-            nn.Upsample(scale_factor=self.upfactor, mode='bilinear', align_corners=True))
-
-        self.init_params()
-
-    def forward(self, x):
-        x = self.adapt_conv(x)
-        return x
-
-    def init_params(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.ConvTranspose2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm2d):  # NN.Batchnorm2d
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-
-
-
-# ==============================================================================================================
-
-
-class ResidualConv(nn.Module):
-    def __init__(self, inchannels):
-        super(ResidualConv, self).__init__()
-        # NN.BatchNorm2d
-        self.conv = nn.Sequential(
-            # nn.BatchNorm2d(num_features=inchannels),
-            nn.ReLU(inplace=False),
-            # nn.Conv2d(in_channels=inchannels, out_channels=inchannels, kernel_size=3, padding=1, stride=1, groups=inchannels,bias=True),
-            # nn.Conv2d(in_channels=inchannels, out_channels=inchannels, kernel_size=1, padding=0, stride=1, groups=1,bias=True)
-            nn.Conv2d(in_channels=inchannels, out_channels=inchannels / 2, kernel_size=3, padding=1, stride=1,
-                      bias=False),
-            nn.BatchNorm2d(num_features=inchannels / 2),
-            nn.ReLU(inplace=False),
-            nn.Conv2d(in_channels=inchannels / 2, out_channels=inchannels, kernel_size=3, padding=1, stride=1,
-                      bias=False)
-        )
-        self.init_params()
-
-    def forward(self, x):
-        x = self.conv(x) + x
-        return x
-
-    def init_params(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.ConvTranspose2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm2d):  # NN.BatchNorm2d
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-
-
-class FeatureFusion(nn.Module):
-    def __init__(self, inchannels, outchannels):
-        super(FeatureFusion, self).__init__()
-        self.conv = ResidualConv(inchannels=inchannels)
-        # NN.BatchNorm2d
-        self.up = nn.Sequential(ResidualConv(inchannels=inchannels),
-                                nn.ConvTranspose2d(in_channels=inchannels, out_channels=outchannels, kernel_size=3,
-                                                   stride=2, padding=1, output_padding=1),
-                                nn.BatchNorm2d(num_features=outchannels),
-                                nn.ReLU(inplace=True))
-
-    def forward(self, lowfeat, highfeat):
-        return self.up(highfeat + self.conv(lowfeat))
-
-    def init_params(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.ConvTranspose2d):
-                # init.kaiming_normal_(m.weight, mode='fan_out')
-                init.normal_(m.weight, std=0.01)
-                # init.xavier_normal_(m.weight)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm2d):  # NN.BatchNorm2d
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                init.normal_(m.weight, std=0.01)
-                if m.bias is not None:
-                    init.constant_(m.bias, 0)
-
-
-class SenceUnderstand(nn.Module):
-    def __init__(self, channels):
-        super(SenceUnderstand, self).__init__()
-        self.channels = channels
-        self.conv1 = nn.Sequential(nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
-                                   nn.ReLU(inplace=True))
-        self.pool = nn.AdaptiveAvgPool2d(8)
-        self.fc = nn.Sequential(nn.Linear(512 * 8 * 8, self.channels),
-                                nn.ReLU(inplace=True))
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_channels=self.channels, out_channels=self.channels, kernel_size=1, padding=0),
-            nn.ReLU(inplace=True))
-        self.initial_params()
-
-    def forward(self, x):
-        n, c, h, w = x.size()
-        x = self.conv1(x)
-        x = self.pool(x)
-        x = x.view(n, -1)
-        x = self.fc(x)
-        x = x.view(n, self.channels, 1, 1)
-        x = self.conv2(x)
-        x = x.repeat(1, 1, h, w)
-        return x
-
-    def initial_params(self, dev=0.01):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                # print torch.sum(m.weight)
-                m.weight.data.normal_(0, dev)
-                if m.bias is not None:
-                    m.bias.data.fill_(0)
-            elif isinstance(m, nn.ConvTranspose2d):
-                # print torch.sum(m.weight)
-                m.weight.data.normal_(0, dev)
-                if m.bias is not None:
-                    m.bias.data.fill_(0)
-            elif isinstance(m, nn.Linear):
-                m.weight.data.normal_(0, dev)
-
-
-if __name__ == '__main__':
-    net = DepthNet(depth=50, pretrained=True)
-    print(net)
-    inputs = torch.ones(4,3,128,128)
-    out = net(inputs)
-    print(out.size())
-

controlnet_aux_local/leres/pix2pix/models/__init__.py

@@ -1,67 +0,0 @@
-"""This package contains modules related to objective functions, optimizations, and network architectures.
-
-To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
-You need to implement the following five functions:
-    -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
-    -- <set_input>:                     unpack data from dataset and apply preprocessing.
-    -- <forward>:                       produce intermediate results.
-    -- <optimize_parameters>:           calculate loss, gradients, and update network weights.
-    -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
-
-In the function <__init__>, you need to define four lists:
-    -- self.loss_names (str list):          specify the training losses that you want to plot and save.
-    -- self.model_names (str list):         define networks used in our training.
-    -- self.visual_names (str list):        specify the images that you want to display and save.
-    -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
-
-Now you can use the model class by specifying flag '--model dummy'.
-See our template model class 'template_model.py' for more details.
-"""
-
-import importlib
-from .base_model import BaseModel
-
-
-def find_model_using_name(model_name):
-    """Import the module "models/[model_name]_model.py".
-
-    In the file, the class called DatasetNameModel() will
-    be instantiated. It has to be a subclass of BaseModel,
-    and it is case-insensitive.
-    """
-    model_filename = "controlnet_aux.leres.pix2pix.models." + model_name + "_model"
-    modellib = importlib.import_module(model_filename)
-    model = None
-    target_model_name = model_name.replace('_', '') + 'model'
-    for name, cls in modellib.__dict__.items():
-        if name.lower() == target_model_name.lower() \
-           and issubclass(cls, BaseModel):
-            model = cls
-
-    if model is None:
-        print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
-        exit(0)
-
-    return model
-
-
-def get_option_setter(model_name):
-    """Return the static method <modify_commandline_options> of the model class."""
-    model_class = find_model_using_name(model_name)
-    return model_class.modify_commandline_options
-
-
-def create_model(opt):
-    """Create a model given the option.
-
-    This function warps the class CustomDatasetDataLoader.
-    This is the main interface between this package and 'train.py'/'test.py'
-
-    Example:
-        >>> from models import create_model
-        >>> model = create_model(opt)
-    """
-    model = find_model_using_name(opt.model)
-    instance = model(opt)
-    print("model [%s] was created" % type(instance).__name__)
-    return instance

controlnet_aux_local/leres/pix2pix/models/base_model.py

@@ -1,244 +0,0 @@
-import gc
-import os
-from abc import ABC, abstractmethod
-from collections import OrderedDict
-
-import torch
-
-from ....util import torch_gc
-from . import networks
-
-
-class BaseModel(ABC):
-    """This class is an abstract base class (ABC) for models.
-    To create a subclass, you need to implement the following five functions:
-        -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
-        -- <set_input>:                     unpack data from dataset and apply preprocessing.
-        -- <forward>:                       produce intermediate results.
-        -- <optimize_parameters>:           calculate losses, gradients, and update network weights.
-        -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
-    """
-
-    def __init__(self, opt):
-        """Initialize the BaseModel class.
-
-        Parameters:
-            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
-
-        When creating your custom class, you need to implement your own initialization.
-        In this function, you should first call <BaseModel.__init__(self, opt)>
-        Then, you need to define four lists:
-            -- self.loss_names (str list):          specify the training losses that you want to plot and save.
-            -- self.model_names (str list):         define networks used in our training.
-            -- self.visual_names (str list):        specify the images that you want to display and save.
-            -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
-        """
-        self.opt = opt
-        self.gpu_ids = opt.gpu_ids
-        self.isTrain = opt.isTrain
-        self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')  # get device name: CPU or GPU
-        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)  # save all the checkpoints to save_dir
-        if opt.preprocess != 'scale_width':  # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
-            torch.backends.cudnn.benchmark = True
-        self.loss_names = []
-        self.model_names = []
-        self.visual_names = []
-        self.optimizers = []
-        self.image_paths = []
-        self.metric = 0  # used for learning rate policy 'plateau'
-
-    @staticmethod
-    def modify_commandline_options(parser, is_train):
-        """Add new model-specific options, and rewrite default values for existing options.
-
-        Parameters:
-            parser          -- original option parser
-            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
-
-        Returns:
-            the modified parser.
-        """
-        return parser
-
-    @abstractmethod
-    def set_input(self, input):
-        """Unpack input data from the dataloader and perform necessary pre-processing steps.
-
-        Parameters:
-            input (dict): includes the data itself and its metadata information.
-        """
-        pass
-
-    @abstractmethod
-    def forward(self):
-        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
-        pass
-
-    @abstractmethod
-    def optimize_parameters(self):
-        """Calculate losses, gradients, and update network weights; called in every training iteration"""
-        pass
-
-    def setup(self, opt):
-        """Load and print networks; create schedulers
-
-        Parameters:
-            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
-        """
-        if self.isTrain:
-            self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
-        if not self.isTrain or opt.continue_train:
-            load_suffix = 'iter_%d' % opt.load_iter if opt.load_iter > 0 else opt.epoch
-            self.load_networks(load_suffix)
-        self.print_networks(opt.verbose)
-
-    def eval(self):
-        """Make models eval mode during test time"""
-        for name in self.model_names:
-            if isinstance(name, str):
-                net = getattr(self, 'net' + name)
-                net.eval()
-
-    def test(self):
-        """Forward function used in test time.
-
-        This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
-        It also calls <compute_visuals> to produce additional visualization results
-        """
-        with torch.no_grad():
-            self.forward()
-            self.compute_visuals()
-
-    def compute_visuals(self):
-        """Calculate additional output images for visdom and HTML visualization"""
-        pass
-
-    def get_image_paths(self):
-        """ Return image paths that are used to load current data"""
-        return self.image_paths
-
-    def update_learning_rate(self):
-        """Update learning rates for all the networks; called at the end of every epoch"""
-        old_lr = self.optimizers[0].param_groups[0]['lr']
-        for scheduler in self.schedulers:
-            if self.opt.lr_policy == 'plateau':
-                scheduler.step(self.metric)
-            else:
-                scheduler.step()
-
-        lr = self.optimizers[0].param_groups[0]['lr']
-        print('learning rate %.7f -> %.7f' % (old_lr, lr))
-
-    def get_current_visuals(self):
-        """Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
-        visual_ret = OrderedDict()
-        for name in self.visual_names:
-            if isinstance(name, str):
-                visual_ret[name] = getattr(self, name)
-        return visual_ret
-
-    def get_current_losses(self):
-        """Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
-        errors_ret = OrderedDict()
-        for name in self.loss_names:
-            if isinstance(name, str):
-                errors_ret[name] = float(getattr(self, 'loss_' + name))  # float(...) works for both scalar tensor and float number
-        return errors_ret
-
-    def save_networks(self, epoch):
-        """Save all the networks to the disk.
-
-        Parameters:
-            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
-        """
-        for name in self.model_names:
-            if isinstance(name, str):
-                save_filename = '%s_net_%s.pth' % (epoch, name)
-                save_path = os.path.join(self.save_dir, save_filename)
-                net = getattr(self, 'net' + name)
-
-                if len(self.gpu_ids) > 0 and torch.cuda.is_available():
-                    torch.save(net.module.cpu().state_dict(), save_path)
-                    net.cuda(self.gpu_ids[0])
-                else:
-                    torch.save(net.cpu().state_dict(), save_path)
-
-    def unload_network(self, name):
-        """Unload network and gc.
-        """
-        if isinstance(name, str):
-            net = getattr(self, 'net' + name)
-            del net
-            gc.collect()
-            torch_gc()
-            return None
-
-    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
-        """Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
-        key = keys[i]
-        if i + 1 == len(keys):  # at the end, pointing to a parameter/buffer
-            if module.__class__.__name__.startswith('InstanceNorm') and \
-                    (key == 'running_mean' or key == 'running_var'):
-                if getattr(module, key) is None:
-                    state_dict.pop('.'.join(keys))
-            if module.__class__.__name__.startswith('InstanceNorm') and \
-               (key == 'num_batches_tracked'):
-                state_dict.pop('.'.join(keys))
-        else:
-            self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
-
-    def load_networks(self, epoch):
-        """Load all the networks from the disk.
-
-        Parameters:
-            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
-        """
-        for name in self.model_names:
-            if isinstance(name, str):
-                load_filename = '%s_net_%s.pth' % (epoch, name)
-                load_path = os.path.join(self.save_dir, load_filename)
-                net = getattr(self, 'net' + name)
-                if isinstance(net, torch.nn.DataParallel):
-                    net = net.module
-                # print('Loading depth boost model from %s' % load_path)
-                # if you are using PyTorch newer than 0.4 (e.g., built from
-                # GitHub source), you can remove str() on self.device
-                state_dict = torch.load(load_path, map_location=str(self.device))
-                if hasattr(state_dict, '_metadata'):
-                    del state_dict._metadata
-
-                # patch InstanceNorm checkpoints prior to 0.4
-                for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
-                    self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
-                net.load_state_dict(state_dict)
-
-    def print_networks(self, verbose):
-        """Print the total number of parameters in the network and (if verbose) network architecture
-
-        Parameters:
-            verbose (bool) -- if verbose: print the network architecture
-        """
-        print('---------- Networks initialized -------------')
-        for name in self.model_names:
-            if isinstance(name, str):
-                net = getattr(self, 'net' + name)
-                num_params = 0
-                for param in net.parameters():
-                    num_params += param.numel()
-                if verbose:
-                    print(net)
-                print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
-        print('-----------------------------------------------')
-
-    def set_requires_grad(self, nets, requires_grad=False):
-        """Set requies_grad=Fasle for all the networks to avoid unnecessary computations
-        Parameters:
-            nets (network list)   -- a list of networks
-            requires_grad (bool)  -- whether the networks require gradients or not
-        """
-        if not isinstance(nets, list):
-            nets = [nets]
-        for net in nets:
-            if net is not None:
-                for param in net.parameters():
-                    param.requires_grad = requires_grad

controlnet_aux_local/leres/pix2pix/models/base_model_hg.py

@@ -1,58 +0,0 @@
-import os
-import torch
-
-class BaseModelHG():
-    def name(self):
-        return 'BaseModel'
-
-    def initialize(self, opt):
-        self.opt = opt
-        self.gpu_ids = opt.gpu_ids
-        self.isTrain = opt.isTrain
-        self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
-        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
-
-    def set_input(self, input):
-        self.input = input
-
-    def forward(self):
-        pass
-
-    # used in test time, no backprop
-    def test(self):
-        pass
-
-    def get_image_paths(self):
-        pass
-
-    def optimize_parameters(self):
-        pass
-
-    def get_current_visuals(self):
-        return self.input
-
-    def get_current_errors(self):
-        return {}
-
-    def save(self, label):
-        pass
-
-    # helper saving function that can be used by subclasses
-    def save_network(self, network, network_label, epoch_label, gpu_ids):
-        save_filename = '_%s_net_%s.pth' % (epoch_label, network_label)
-        save_path = os.path.join(self.save_dir, save_filename)
-        torch.save(network.cpu().state_dict(), save_path)
-        if len(gpu_ids) and torch.cuda.is_available():
-            network.cuda(device_id=gpu_ids[0])
-
-    # helper loading function that can be used by subclasses
-    def load_network(self, network, network_label, epoch_label):
-        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
-        save_path = os.path.join(self.save_dir, save_filename)
-        print(save_path)
-        model = torch.load(save_path)
-        return model
-        # network.load_state_dict(torch.load(save_path))
-
-    def update_learning_rate():
-        pass

controlnet_aux_local/leres/pix2pix/models/networks.py

@@ -1,623 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.nn import init
-import functools
-from torch.optim import lr_scheduler
-
-
-###############################################################################
-# Helper Functions
-###############################################################################
-
-
-class Identity(nn.Module):
-    def forward(self, x):
-        return x
-
-
-def get_norm_layer(norm_type='instance'):
-    """Return a normalization layer
-
-    Parameters:
-        norm_type (str) -- the name of the normalization layer: batch | instance | none
-
-    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
-    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
-    """
-    if norm_type == 'batch':
-        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
-    elif norm_type == 'instance':
-        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
-    elif norm_type == 'none':
-        def norm_layer(x): return Identity()
-    else:
-        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
-    return norm_layer
-
-
-def get_scheduler(optimizer, opt):
-    """Return a learning rate scheduler
-
-    Parameters:
-        optimizer          -- the optimizer of the network
-        opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　
-                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
-
-    For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs
-    and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.
-    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
-    See https://pytorch.org/docs/stable/optim.html for more details.
-    """
-    if opt.lr_policy == 'linear':
-        def lambda_rule(epoch):
-            lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)
-            return lr_l
-        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
-    elif opt.lr_policy == 'step':
-        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
-    elif opt.lr_policy == 'plateau':
-        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
-    elif opt.lr_policy == 'cosine':
-        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)
-    else:
-        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
-    return scheduler
-
-
-def init_weights(net, init_type='normal', init_gain=0.02):
-    """Initialize network weights.
-
-    Parameters:
-        net (network)   -- network to be initialized
-        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
-        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.
-
-    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
-    work better for some applications. Feel free to try yourself.
-    """
-    def init_func(m):  # define the initialization function
-        classname = m.__class__.__name__
-        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
-            if init_type == 'normal':
-                init.normal_(m.weight.data, 0.0, init_gain)
-            elif init_type == 'xavier':
-                init.xavier_normal_(m.weight.data, gain=init_gain)
-            elif init_type == 'kaiming':
-                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
-            elif init_type == 'orthogonal':
-                init.orthogonal_(m.weight.data, gain=init_gain)
-            else:
-                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
-            if hasattr(m, 'bias') and m.bias is not None:
-                init.constant_(m.bias.data, 0.0)
-        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
-            init.normal_(m.weight.data, 1.0, init_gain)
-            init.constant_(m.bias.data, 0.0)
-
-    # print('initialize network with %s' % init_type)
-    net.apply(init_func)  # apply the initialization function <init_func>
-
-
-def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
-    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
-    Parameters:
-        net (network)      -- the network to be initialized
-        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
-        gain (float)       -- scaling factor for normal, xavier and orthogonal.
-        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
-
-    Return an initialized network.
-    """
-    if len(gpu_ids) > 0:
-        assert(torch.cuda.is_available())
-        net.to(gpu_ids[0])
-        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
-    init_weights(net, init_type, init_gain=init_gain)
-    return net
-
-
-def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
-    """Create a generator
-
-    Parameters:
-        input_nc (int) -- the number of channels in input images
-        output_nc (int) -- the number of channels in output images
-        ngf (int) -- the number of filters in the last conv layer
-        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
-        norm (str) -- the name of normalization layers used in the network: batch | instance | none
-        use_dropout (bool) -- if use dropout layers.
-        init_type (str)    -- the name of our initialization method.
-        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
-        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
-
-    Returns a generator
-
-    Our current implementation provides two types of generators:
-        U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
-        The original U-Net paper: https://arxiv.org/abs/1505.04597
-
-        Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
-        Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
-        We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).
-
-
-    The generator has been initialized by <init_net>. It uses RELU for non-linearity.
-    """
-    net = None
-    norm_layer = get_norm_layer(norm_type=norm)
-
-    if netG == 'resnet_9blocks':
-        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
-    elif netG == 'resnet_6blocks':
-        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
-    elif netG == 'resnet_12blocks':
-        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=12)
-    elif netG == 'unet_128':
-        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
-    elif netG == 'unet_256':
-        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
-    elif netG == 'unet_672':
-        net = UnetGenerator(input_nc, output_nc, 5, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
-    elif netG == 'unet_960':
-        net = UnetGenerator(input_nc, output_nc, 6, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
-    elif netG == 'unet_1024':
-        net = UnetGenerator(input_nc, output_nc, 10, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
-    else:
-        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
-    return init_net(net, init_type, init_gain, gpu_ids)
-
-
-def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]):
-    """Create a discriminator
-
-    Parameters:
-        input_nc (int)     -- the number of channels in input images
-        ndf (int)          -- the number of filters in the first conv layer
-        netD (str)         -- the architecture's name: basic | n_layers | pixel
-        n_layers_D (int)   -- the number of conv layers in the discriminator; effective when netD=='n_layers'
-        norm (str)         -- the type of normalization layers used in the network.
-        init_type (str)    -- the name of the initialization method.
-        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
-        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
-
-    Returns a discriminator
-
-    Our current implementation provides three types of discriminators:
-        [basic]: 'PatchGAN' classifier described in the original pix2pix paper.
-        It can classify whether 70×70 overlapping patches are real or fake.
-        Such a patch-level discriminator architecture has fewer parameters
-        than a full-image discriminator and can work on arbitrarily-sized images
-        in a fully convolutional fashion.
-
-        [n_layers]: With this mode, you can specify the number of conv layers in the discriminator
-        with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)
-
-        [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
-        It encourages greater color diversity but has no effect on spatial statistics.
-
-    The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
-    """
-    net = None
-    norm_layer = get_norm_layer(norm_type=norm)
-
-    if netD == 'basic':  # default PatchGAN classifier
-        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
-    elif netD == 'n_layers':  # more options
-        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)
-    elif netD == 'pixel':     # classify if each pixel is real or fake
-        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
-    else:
-        raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
-    return init_net(net, init_type, init_gain, gpu_ids)
-
-
-##############################################################################
-# Classes
-##############################################################################
-class GANLoss(nn.Module):
-    """Define different GAN objectives.
-
-    The GANLoss class abstracts away the need to create the target label tensor
-    that has the same size as the input.
-    """
-
-    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
-        """ Initialize the GANLoss class.
-
-        Parameters:
-            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
-            target_real_label (bool) - - label for a real image
-            target_fake_label (bool) - - label of a fake image
-
-        Note: Do not use sigmoid as the last layer of Discriminator.
-        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
-        """
-        super(GANLoss, self).__init__()
-        self.register_buffer('real_label', torch.tensor(target_real_label))
-        self.register_buffer('fake_label', torch.tensor(target_fake_label))
-        self.gan_mode = gan_mode
-        if gan_mode == 'lsgan':
-            self.loss = nn.MSELoss()
-        elif gan_mode == 'vanilla':
-            self.loss = nn.BCEWithLogitsLoss()
-        elif gan_mode in ['wgangp']:
-            self.loss = None
-        else:
-            raise NotImplementedError('gan mode %s not implemented' % gan_mode)
-
-    def get_target_tensor(self, prediction, target_is_real):
-        """Create label tensors with the same size as the input.
-
-        Parameters:
-            prediction (tensor) - - tpyically the prediction from a discriminator
-            target_is_real (bool) - - if the ground truth label is for real images or fake images
-
-        Returns:
-            A label tensor filled with ground truth label, and with the size of the input
-        """
-
-        if target_is_real:
-            target_tensor = self.real_label
-        else:
-            target_tensor = self.fake_label
-        return target_tensor.expand_as(prediction)
-
-    def __call__(self, prediction, target_is_real):
-        """Calculate loss given Discriminator's output and grount truth labels.
-
-        Parameters:
-            prediction (tensor) - - tpyically the prediction output from a discriminator
-            target_is_real (bool) - - if the ground truth label is for real images or fake images
-
-        Returns:
-            the calculated loss.
-        """
-        if self.gan_mode in ['lsgan', 'vanilla']:
-            target_tensor = self.get_target_tensor(prediction, target_is_real)
-            loss = self.loss(prediction, target_tensor)
-        elif self.gan_mode == 'wgangp':
-            if target_is_real:
-                loss = -prediction.mean()
-            else:
-                loss = prediction.mean()
-        return loss
-
-
-def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):
-    """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
-
-    Arguments:
-        netD (network)              -- discriminator network
-        real_data (tensor array)    -- real images
-        fake_data (tensor array)    -- generated images from the generator
-        device (str)                -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
-        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].
-        constant (float)            -- the constant used in formula ( ||gradient||_2 - constant)^2
-        lambda_gp (float)           -- weight for this loss
-
-    Returns the gradient penalty loss
-    """
-    if lambda_gp > 0.0:
-        if type == 'real':   # either use real images, fake images, or a linear interpolation of two.
-            interpolatesv = real_data
-        elif type == 'fake':
-            interpolatesv = fake_data
-        elif type == 'mixed':
-            alpha = torch.rand(real_data.shape[0], 1, device=device)
-            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
-            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
-        else:
-            raise NotImplementedError('{} not implemented'.format(type))
-        interpolatesv.requires_grad_(True)
-        disc_interpolates = netD(interpolatesv)
-        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,
-                                        grad_outputs=torch.ones(disc_interpolates.size()).to(device),
-                                        create_graph=True, retain_graph=True, only_inputs=True)
-        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data
-        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp        # added eps
-        return gradient_penalty, gradients
-    else:
-        return 0.0, None
-
-
-class ResnetGenerator(nn.Module):
-    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
-
-    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
-    """
-
-    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
-        """Construct a Resnet-based generator
-
-        Parameters:
-            input_nc (int)      -- the number of channels in input images
-            output_nc (int)     -- the number of channels in output images
-            ngf (int)           -- the number of filters in the last conv layer
-            norm_layer          -- normalization layer
-            use_dropout (bool)  -- if use dropout layers
-            n_blocks (int)      -- the number of ResNet blocks
-            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
-        """
-        assert(n_blocks >= 0)
-        super(ResnetGenerator, self).__init__()
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-
-        model = [nn.ReflectionPad2d(3),
-                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
-                 norm_layer(ngf),
-                 nn.ReLU(True)]
-
-        n_downsampling = 2
-        for i in range(n_downsampling):  # add downsampling layers
-            mult = 2 ** i
-            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
-                      norm_layer(ngf * mult * 2),
-                      nn.ReLU(True)]
-
-        mult = 2 ** n_downsampling
-        for i in range(n_blocks):       # add ResNet blocks
-
-            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
-
-        for i in range(n_downsampling):  # add upsampling layers
-            mult = 2 ** (n_downsampling - i)
-            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
-                                         kernel_size=3, stride=2,
-                                         padding=1, output_padding=1,
-                                         bias=use_bias),
-                      norm_layer(int(ngf * mult / 2)),
-                      nn.ReLU(True)]
-        model += [nn.ReflectionPad2d(3)]
-        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
-        model += [nn.Tanh()]
-
-        self.model = nn.Sequential(*model)
-
-    def forward(self, input):
-        """Standard forward"""
-        return self.model(input)
-
-
-class ResnetBlock(nn.Module):
-    """Define a Resnet block"""
-
-    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
-        """Initialize the Resnet block
-
-        A resnet block is a conv block with skip connections
-        We construct a conv block with build_conv_block function,
-        and implement skip connections in <forward> function.
-        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
-        """
-        super(ResnetBlock, self).__init__()
-        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
-
-    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
-        """Construct a convolutional block.
-
-        Parameters:
-            dim (int)           -- the number of channels in the conv layer.
-            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
-            norm_layer          -- normalization layer
-            use_dropout (bool)  -- if use dropout layers.
-            use_bias (bool)     -- if the conv layer uses bias or not
-
-        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
-        """
-        conv_block = []
-        p = 0
-        if padding_type == 'reflect':
-            conv_block += [nn.ReflectionPad2d(1)]
-        elif padding_type == 'replicate':
-            conv_block += [nn.ReplicationPad2d(1)]
-        elif padding_type == 'zero':
-            p = 1
-        else:
-            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
-
-        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
-        if use_dropout:
-            conv_block += [nn.Dropout(0.5)]
-
-        p = 0
-        if padding_type == 'reflect':
-            conv_block += [nn.ReflectionPad2d(1)]
-        elif padding_type == 'replicate':
-            conv_block += [nn.ReplicationPad2d(1)]
-        elif padding_type == 'zero':
-            p = 1
-        else:
-            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
-        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
-
-        return nn.Sequential(*conv_block)
-
-    def forward(self, x):
-        """Forward function (with skip connections)"""
-        out = x + self.conv_block(x)  # add skip connections
-        return out
-
-
-class UnetGenerator(nn.Module):
-    """Create a Unet-based generator"""
-
-    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
-        """Construct a Unet generator
-        Parameters:
-            input_nc (int)  -- the number of channels in input images
-            output_nc (int) -- the number of channels in output images
-            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
-                                image of size 128x128 will become of size 1x1 # at the bottleneck
-            ngf (int)       -- the number of filters in the last conv layer
-            norm_layer      -- normalization layer
-
-        We construct the U-Net from the innermost layer to the outermost layer.
-        It is a recursive process.
-        """
-        super(UnetGenerator, self).__init__()
-        # construct unet structure
-        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer
-        for i in range(num_downs - 5):          # add intermediate layers with ngf * 8 filters
-            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
-        # gradually reduce the number of filters from ngf * 8 to ngf
-        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer
-
-    def forward(self, input):
-        """Standard forward"""
-        return self.model(input)
-
-
-class UnetSkipConnectionBlock(nn.Module):
-    """Defines the Unet submodule with skip connection.
-        X -------------------identity----------------------
-        |-- downsampling -- |submodule| -- upsampling --|
-    """
-
-    def __init__(self, outer_nc, inner_nc, input_nc=None,
-                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
-        """Construct a Unet submodule with skip connections.
-
-        Parameters:
-            outer_nc (int) -- the number of filters in the outer conv layer
-            inner_nc (int) -- the number of filters in the inner conv layer
-            input_nc (int) -- the number of channels in input images/features
-            submodule (UnetSkipConnectionBlock) -- previously defined submodules
-            outermost (bool)    -- if this module is the outermost module
-            innermost (bool)    -- if this module is the innermost module
-            norm_layer          -- normalization layer
-            use_dropout (bool)  -- if use dropout layers.
-        """
-        super(UnetSkipConnectionBlock, self).__init__()
-        self.outermost = outermost
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-        if input_nc is None:
-            input_nc = outer_nc
-        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
-                             stride=2, padding=1, bias=use_bias)
-        downrelu = nn.LeakyReLU(0.2, True)
-        downnorm = norm_layer(inner_nc)
-        uprelu = nn.ReLU(True)
-        upnorm = norm_layer(outer_nc)
-
-        if outermost:
-            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1)
-            down = [downconv]
-            up = [uprelu, upconv, nn.Tanh()]
-            model = down + [submodule] + up
-        elif innermost:
-            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1, bias=use_bias)
-            down = [downrelu, downconv]
-            up = [uprelu, upconv, upnorm]
-            model = down + up
-        else:
-            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1, bias=use_bias)
-            down = [downrelu, downconv, downnorm]
-            up = [uprelu, upconv, upnorm]
-
-            if use_dropout:
-                model = down + [submodule] + up + [nn.Dropout(0.5)]
-            else:
-                model = down + [submodule] + up
-
-        self.model = nn.Sequential(*model)
-
-    def forward(self, x):
-        if self.outermost:
-            return self.model(x)
-        else:   # add skip connections
-            return torch.cat([x, self.model(x)], 1)
-
-
-class NLayerDiscriminator(nn.Module):
-    """Defines a PatchGAN discriminator"""
-
-    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
-        """Construct a PatchGAN discriminator
-
-        Parameters:
-            input_nc (int)  -- the number of channels in input images
-            ndf (int)       -- the number of filters in the last conv layer
-            n_layers (int)  -- the number of conv layers in the discriminator
-            norm_layer      -- normalization layer
-        """
-        super(NLayerDiscriminator, self).__init__()
-        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-
-        kw = 4
-        padw = 1
-        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
-        nf_mult = 1
-        nf_mult_prev = 1
-        for n in range(1, n_layers):  # gradually increase the number of filters
-            nf_mult_prev = nf_mult
-            nf_mult = min(2 ** n, 8)
-            sequence += [
-                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
-                norm_layer(ndf * nf_mult),
-                nn.LeakyReLU(0.2, True)
-            ]
-
-        nf_mult_prev = nf_mult
-        nf_mult = min(2 ** n_layers, 8)
-        sequence += [
-            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
-            norm_layer(ndf * nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-
-        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
-        self.model = nn.Sequential(*sequence)
-
-    def forward(self, input):
-        """Standard forward."""
-        return self.model(input)
-
-
-class PixelDiscriminator(nn.Module):
-    """Defines a 1x1 PatchGAN discriminator (pixelGAN)"""
-
-    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):
-        """Construct a 1x1 PatchGAN discriminator
-
-        Parameters:
-            input_nc (int)  -- the number of channels in input images
-            ndf (int)       -- the number of filters in the last conv layer
-            norm_layer      -- normalization layer
-        """
-        super(PixelDiscriminator, self).__init__()
-        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-
-        self.net = [
-            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
-            nn.LeakyReLU(0.2, True),
-            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
-            norm_layer(ndf * 2),
-            nn.LeakyReLU(0.2, True),
-            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
-
-        self.net = nn.Sequential(*self.net)
-
-    def forward(self, input):
-        """Standard forward."""
-        return self.net(input)

controlnet_aux_local/leres/pix2pix/models/pix2pix4depth_model.py

@@ -1,155 +0,0 @@
-import torch
-from .base_model import BaseModel
-from . import networks
-
-
-class Pix2Pix4DepthModel(BaseModel):
-    """ This class implements the pix2pix model, for learning a mapping from input images to output images given paired data.
-
-    The model training requires '--dataset_mode aligned' dataset.
-    By default, it uses a '--netG unet256' U-Net generator,
-    a '--netD basic' discriminator (PatchGAN),
-    and a '--gan_mode' vanilla GAN loss (the cross-entropy objective used in the orignal GAN paper).
-
-    pix2pix paper: https://arxiv.org/pdf/1611.07004.pdf
-    """
-    @staticmethod
-    def modify_commandline_options(parser, is_train=True):
-        """Add new dataset-specific options, and rewrite default values for existing options.
-
-        Parameters:
-            parser          -- original option parser
-            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
-
-        Returns:
-            the modified parser.
-
-        For pix2pix, we do not use image buffer
-        The training objective is: GAN Loss + lambda_L1 * ||G(A)-B||_1
-        By default, we use vanilla GAN loss, UNet with batchnorm, and aligned datasets.
-        """
-        # changing the default values to match the pix2pix paper (https://phillipi.github.io/pix2pix/)
-        parser.set_defaults(input_nc=2,output_nc=1,norm='none', netG='unet_1024', dataset_mode='depthmerge')
-        if is_train:
-            parser.set_defaults(pool_size=0, gan_mode='vanilla',)
-            parser.add_argument('--lambda_L1', type=float, default=1000, help='weight for L1 loss')
-        return parser
-
-    def __init__(self, opt):
-        """Initialize the pix2pix class.
-
-        Parameters:
-            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
-        """
-        BaseModel.__init__(self, opt)
-        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
-
-        self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake']
-        # self.loss_names = ['G_L1']
-
-        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
-        if self.isTrain:
-            self.visual_names = ['outer','inner', 'fake_B', 'real_B']
-        else:
-            self.visual_names = ['fake_B']
-
-        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
-        if self.isTrain:
-            self.model_names = ['G','D']
-        else:  # during test time, only load G
-            self.model_names = ['G']
-
-        # define networks (both generator and discriminator)
-        self.netG = networks.define_G(opt.input_nc, opt.output_nc, 64, 'unet_1024', 'none',
-                                      False, 'normal', 0.02, self.gpu_ids)
-
-        if self.isTrain:  # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
-            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
-                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
-
-        if self.isTrain:
-            # define loss functions
-            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
-            self.criterionL1 = torch.nn.L1Loss()
-            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
-            self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=1e-4, betas=(opt.beta1, 0.999))
-            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=2e-06, betas=(opt.beta1, 0.999))
-            self.optimizers.append(self.optimizer_G)
-            self.optimizers.append(self.optimizer_D)
-
-    def set_input_train(self, input):
-        self.outer = input['data_outer'].to(self.device)
-        self.outer = torch.nn.functional.interpolate(self.outer,(1024,1024),mode='bilinear',align_corners=False)
-
-        self.inner = input['data_inner'].to(self.device)
-        self.inner = torch.nn.functional.interpolate(self.inner,(1024,1024),mode='bilinear',align_corners=False)
-
-        self.image_paths = input['image_path']
-
-        if self.isTrain:
-            self.gtfake = input['data_gtfake'].to(self.device)
-            self.gtfake = torch.nn.functional.interpolate(self.gtfake, (1024, 1024), mode='bilinear', align_corners=False)
-            self.real_B = self.gtfake
-
-        self.real_A = torch.cat((self.outer, self.inner), 1)
-
-    def set_input(self, outer, inner):
-        inner = torch.from_numpy(inner).unsqueeze(0).unsqueeze(0)
-        outer = torch.from_numpy(outer).unsqueeze(0).unsqueeze(0)
-
-        inner = (inner - torch.min(inner))/(torch.max(inner)-torch.min(inner))
-        outer = (outer - torch.min(outer))/(torch.max(outer)-torch.min(outer))
-
-        inner = self.normalize(inner)
-        outer = self.normalize(outer)
-
-        self.real_A = torch.cat((outer, inner), 1).to(self.device)
-
-
-    def normalize(self, input):
-        input = input * 2
-        input = input - 1
-        return input
-
-    def forward(self):
-        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
-        self.fake_B = self.netG(self.real_A)  # G(A)
-
-    def backward_D(self):
-        """Calculate GAN loss for the discriminator"""
-        # Fake; stop backprop to the generator by detaching fake_B
-        fake_AB = torch.cat((self.real_A, self.fake_B), 1)  # we use conditional GANs; we need to feed both input and output to the discriminator
-        pred_fake = self.netD(fake_AB.detach())
-        self.loss_D_fake = self.criterionGAN(pred_fake, False)
-        # Real
-        real_AB = torch.cat((self.real_A, self.real_B), 1)
-        pred_real = self.netD(real_AB)
-        self.loss_D_real = self.criterionGAN(pred_real, True)
-        # combine loss and calculate gradients
-        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
-        self.loss_D.backward()
-
-    def backward_G(self):
-        """Calculate GAN and L1 loss for the generator"""
-        # First, G(A) should fake the discriminator
-        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
-        pred_fake = self.netD(fake_AB)
-        self.loss_G_GAN = self.criterionGAN(pred_fake, True)
-        # Second, G(A) = B
-        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1
-        # combine loss and calculate gradients
-        self.loss_G = self.loss_G_L1 + self.loss_G_GAN
-        self.loss_G.backward()
-
-    def optimize_parameters(self):
-        self.forward()                   # compute fake images: G(A)
-        # update D
-        self.set_requires_grad(self.netD, True)  # enable backprop for D
-        self.optimizer_D.zero_grad()     # set D's gradients to zero
-        self.backward_D()                # calculate gradients for D
-        self.optimizer_D.step()          # update D's weights
-        # update G
-        self.set_requires_grad(self.netD, False)  # D requires no gradients when optimizing G
-        self.optimizer_G.zero_grad()        # set G's gradients to zero
-        self.backward_G()                   # calculate graidents for G
-        self.optimizer_G.step()             # udpate G's weights

controlnet_aux_local/leres/pix2pix/options/__init__.py

@@ -1 +0,0 @@
-"""This package options includes option modules: training options, test options, and basic options (used in both training and test)."""

controlnet_aux_local/leres/pix2pix/options/base_options.py

@@ -1,156 +0,0 @@
-import argparse
-import os
-from ...pix2pix.util import util
-# import torch
-from ...pix2pix import models
-# import pix2pix.data
-import numpy as np
-
-class BaseOptions():
-    """This class defines options used during both training and test time.
-
-    It also implements several helper functions such as parsing, printing, and saving the options.
-    It also gathers additional options defined in <modify_commandline_options> functions in both dataset class and model class.
-    """
-
-    def __init__(self):
-        """Reset the class; indicates the class hasn't been initailized"""
-        self.initialized = False
-
-    def initialize(self, parser):
-        """Define the common options that are used in both training and test."""
-        # basic parameters
-        parser.add_argument('--dataroot', help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')
-        parser.add_argument('--name', type=str, default='void', help='mahdi_unet_new, scaled_unet')
-        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
-        parser.add_argument('--checkpoints_dir', type=str, default='./pix2pix/checkpoints', help='models are saved here')
-        # model parameters
-        parser.add_argument('--model', type=str, default='cycle_gan', help='chooses which model to use. [cycle_gan | pix2pix | test | colorization]')
-        parser.add_argument('--input_nc', type=int, default=2, help='# of input image channels: 3 for RGB and 1 for grayscale')
-        parser.add_argument('--output_nc', type=int, default=1, help='# of output image channels: 3 for RGB and 1 for grayscale')
-        parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in the last conv layer')
-        parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in the first conv layer')
-        parser.add_argument('--netD', type=str, default='basic', help='specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator')
-        parser.add_argument('--netG', type=str, default='resnet_9blocks', help='specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]')
-        parser.add_argument('--n_layers_D', type=int, default=3, help='only used if netD==n_layers')
-        parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization [instance | batch | none]')
-        parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal | xavier | kaiming | orthogonal]')
-        parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
-        parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator')
-        # dataset parameters
-        parser.add_argument('--dataset_mode', type=str, default='unaligned', help='chooses how datasets are loaded. [unaligned | aligned | single | colorization]')
-        parser.add_argument('--direction', type=str, default='AtoB', help='AtoB or BtoA')
-        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
-        parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
-        parser.add_argument('--batch_size', type=int, default=1, help='input batch size')
-        parser.add_argument('--load_size', type=int, default=672, help='scale images to this size')
-        parser.add_argument('--crop_size', type=int, default=672, help='then crop to this size')
-        parser.add_argument('--max_dataset_size', type=int, default=10000, help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
-        parser.add_argument('--preprocess', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]')
-        parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
-        parser.add_argument('--display_winsize', type=int, default=256, help='display window size for both visdom and HTML')
-        # additional parameters
-        parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
-        parser.add_argument('--load_iter', type=int, default='0', help='which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]')
-        parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')
-        parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}')
-
-        parser.add_argument('--data_dir', type=str, required=False,
-                            help='input files directory images can be .png .jpg .tiff')
-        parser.add_argument('--output_dir', type=str, required=False,
-                            help='result dir. result depth will be png. vides are JMPG as avi')
-        parser.add_argument('--savecrops', type=int, required=False)
-        parser.add_argument('--savewholeest', type=int, required=False)
-        parser.add_argument('--output_resolution', type=int, required=False,
-                            help='0 for no restriction 1 for resize to input size')
-        parser.add_argument('--net_receptive_field_size', type=int, required=False)
-        parser.add_argument('--pix2pixsize', type=int, required=False)
-        parser.add_argument('--generatevideo', type=int, required=False)
-        parser.add_argument('--depthNet', type=int, required=False, help='0: midas 1:strurturedRL')
-        parser.add_argument('--R0', action='store_true')
-        parser.add_argument('--R20', action='store_true')
-        parser.add_argument('--Final', action='store_true')
-        parser.add_argument('--colorize_results', action='store_true')
-        parser.add_argument('--max_res', type=float, default=np.inf)
-
-        self.initialized = True
-        return parser
-
-    def gather_options(self):
-        """Initialize our parser with basic options(only once).
-        Add additional model-specific and dataset-specific options.
-        These options are defined in the <modify_commandline_options> function
-        in model and dataset classes.
-        """
-        if not self.initialized:  # check if it has been initialized
-            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-            parser = self.initialize(parser)
-
-        # get the basic options
-        opt, _ = parser.parse_known_args()
-
-        # modify model-related parser options
-        model_name = opt.model
-        model_option_setter = models.get_option_setter(model_name)
-        parser = model_option_setter(parser, self.isTrain)
-        opt, _ = parser.parse_known_args()  # parse again with new defaults
-
-        # modify dataset-related parser options
-        # dataset_name = opt.dataset_mode
-        # dataset_option_setter = pix2pix.data.get_option_setter(dataset_name)
-        # parser = dataset_option_setter(parser, self.isTrain)
-
-        # save and return the parser
-        self.parser = parser
-        #return parser.parse_args() #EVIL
-        return opt
-
-    def print_options(self, opt):
-        """Print and save options
-
-        It will print both current options and default values(if different).
-        It will save options into a text file / [checkpoints_dir] / opt.txt
-        """
-        message = ''
-        message += '----------------- Options ---------------\n'
-        for k, v in sorted(vars(opt).items()):
-            comment = ''
-            default = self.parser.get_default(k)
-            if v != default:
-                comment = '\t[default: %s]' % str(default)
-            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
-        message += '----------------- End -------------------'
-        print(message)
-
-        # save to the disk
-        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
-        util.mkdirs(expr_dir)
-        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
-        with open(file_name, 'wt') as opt_file:
-            opt_file.write(message)
-            opt_file.write('\n')
-
-    def parse(self):
-        """Parse our options, create checkpoints directory suffix, and set up gpu device."""
-        opt = self.gather_options()
-        opt.isTrain = self.isTrain   # train or test
-
-        # process opt.suffix
-        if opt.suffix:
-            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
-            opt.name = opt.name + suffix
-
-        #self.print_options(opt)
-
-        # set gpu ids
-        str_ids = opt.gpu_ids.split(',')
-        opt.gpu_ids = []
-        for str_id in str_ids:
-            id = int(str_id)
-            if id >= 0:
-                opt.gpu_ids.append(id)
-        #if len(opt.gpu_ids) > 0:
-        #    torch.cuda.set_device(opt.gpu_ids[0])
-
-        self.opt = opt
-        return self.opt

controlnet_aux_local/leres/pix2pix/options/test_options.py

@@ -1,22 +0,0 @@
-from .base_options import BaseOptions
-
-
-class TestOptions(BaseOptions):
-    """This class includes test options.
-
-    It also includes shared options defined in BaseOptions.
-    """
-
-    def initialize(self, parser):
-        parser = BaseOptions.initialize(self, parser)  # define shared options
-        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
-        parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
-        # Dropout and Batchnorm has different behavioir during training and test.
-        parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
-        parser.add_argument('--num_test', type=int, default=50, help='how many test images to run')
-        # rewrite devalue values
-        parser.set_defaults(model='pix2pix4depth')
-        # To avoid cropping, the load_size should be the same as crop_size
-        parser.set_defaults(load_size=parser.get_default('crop_size'))
-        self.isTrain = False
-        return parser

controlnet_aux_local/leres/pix2pix/util/__init__.py

@@ -1 +0,0 @@
-"""This package includes a miscellaneous collection of useful helper functions."""

controlnet_aux_local/leres/pix2pix/util/util.py

@@ -1,105 +0,0 @@
-"""This module contains simple helper functions """
-from __future__ import print_function
-import torch
-import numpy as np
-from PIL import Image
-import os
-
-
-def tensor2im(input_image, imtype=np.uint16):
-    """"Converts a Tensor array into a numpy image array.
-
-    Parameters:
-        input_image (tensor) --  the input image tensor array
-        imtype (type)        --  the desired type of the converted numpy array
-    """
-    if not isinstance(input_image, np.ndarray):
-        if isinstance(input_image, torch.Tensor):  # get the data from a variable
-            image_tensor = input_image.data
-        else:
-            return input_image
-        image_numpy = torch.squeeze(image_tensor).cpu().numpy()  # convert it into a numpy array
-        image_numpy = (image_numpy + 1) / 2.0 * (2**16-1) #
-    else:  # if it is a numpy array, do nothing
-        image_numpy = input_image
-    return image_numpy.astype(imtype)
-
-
-def diagnose_network(net, name='network'):
-    """Calculate and print the mean of average absolute(gradients)
-
-    Parameters:
-        net (torch network) -- Torch network
-        name (str) -- the name of the network
-    """
-    mean = 0.0
-    count = 0
-    for param in net.parameters():
-        if param.grad is not None:
-            mean += torch.mean(torch.abs(param.grad.data))
-            count += 1
-    if count > 0:
-        mean = mean / count
-    print(name)
-    print(mean)
-
-
-def save_image(image_numpy, image_path, aspect_ratio=1.0):
-    """Save a numpy image to the disk
-
-    Parameters:
-        image_numpy (numpy array) -- input numpy array
-        image_path (str)          -- the path of the image
-    """
-    image_pil = Image.fromarray(image_numpy)
-
-    image_pil = image_pil.convert('I;16')
-
-    # image_pil = Image.fromarray(image_numpy)
-    # h, w, _ = image_numpy.shape
-    #
-    # if aspect_ratio > 1.0:
-    #     image_pil = image_pil.resize((h, int(w * aspect_ratio)), Image.BICUBIC)
-    # if aspect_ratio < 1.0:
-    #     image_pil = image_pil.resize((int(h / aspect_ratio), w), Image.BICUBIC)
-
-    image_pil.save(image_path)
-
-
-def print_numpy(x, val=True, shp=False):
-    """Print the mean, min, max, median, std, and size of a numpy array
-
-    Parameters:
-        val (bool) -- if print the values of the numpy array
-        shp (bool) -- if print the shape of the numpy array
-    """
-    x = x.astype(np.float64)
-    if shp:
-        print('shape,', x.shape)
-    if val:
-        x = x.flatten()
-        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
-            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
-
-
-def mkdirs(paths):
-    """create empty directories if they don't exist
-
-    Parameters:
-        paths (str list) -- a list of directory paths
-    """
-    if isinstance(paths, list) and not isinstance(paths, str):
-        for path in paths:
-            mkdir(path)
-    else:
-        mkdir(paths)
-
-
-def mkdir(path):
-    """create a single empty directory if it didn't exist
-
-    Parameters:
-        path (str) -- a single directory path
-    """
-    if not os.path.exists(path):
-        os.makedirs(path)

controlnet_aux_local/lineart/__init__.py

@@ -1,167 +0,0 @@
-import os
-import warnings
-
-import cv2
-import numpy as np
-import torch
-import torch.nn as nn
-from einops import rearrange
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from ..util import HWC3, resize_image
-
-norm_layer = nn.InstanceNorm2d
-
-
-class ResidualBlock(nn.Module):
-    def __init__(self, in_features):
-        super(ResidualBlock, self).__init__()
-
-        conv_block = [  nn.ReflectionPad2d(1),
-                        nn.Conv2d(in_features, in_features, 3),
-                        norm_layer(in_features),
-                        nn.ReLU(inplace=True),
-                        nn.ReflectionPad2d(1),
-                        nn.Conv2d(in_features, in_features, 3),
-                        norm_layer(in_features)
-                        ]
-
-        self.conv_block = nn.Sequential(*conv_block)
-
-    def forward(self, x):
-        return x + self.conv_block(x)
-
-
-class Generator(nn.Module):
-    def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):
-        super(Generator, self).__init__()
-
-        # Initial convolution block
-        model0 = [   nn.ReflectionPad2d(3),
-                    nn.Conv2d(input_nc, 64, 7),
-                    norm_layer(64),
-                    nn.ReLU(inplace=True) ]
-        self.model0 = nn.Sequential(*model0)
-
-        # Downsampling
-        model1 = []
-        in_features = 64
-        out_features = in_features*2
-        for _ in range(2):
-            model1 += [  nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
-                        norm_layer(out_features),
-                        nn.ReLU(inplace=True) ]
-            in_features = out_features
-            out_features = in_features*2
-        self.model1 = nn.Sequential(*model1)
-
-        model2 = []
-        # Residual blocks
-        for _ in range(n_residual_blocks):
-            model2 += [ResidualBlock(in_features)]
-        self.model2 = nn.Sequential(*model2)
-
-        # Upsampling
-        model3 = []
-        out_features = in_features//2
-        for _ in range(2):
-            model3 += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
-                        norm_layer(out_features),
-                        nn.ReLU(inplace=True) ]
-            in_features = out_features
-            out_features = in_features//2
-        self.model3 = nn.Sequential(*model3)
-
-        # Output layer
-        model4 = [  nn.ReflectionPad2d(3),
-                        nn.Conv2d(64, output_nc, 7)]
-        if sigmoid:
-            model4 += [nn.Sigmoid()]
-
-        self.model4 = nn.Sequential(*model4)
-
-    def forward(self, x, cond=None):
-        out = self.model0(x)
-        out = self.model1(out)
-        out = self.model2(out)
-        out = self.model3(out)
-        out = self.model4(out)
-
-        return out
-
-
-class LineartDetector:
-    def __init__(self, model, coarse_model):
-        self.model = model
-        self.model_coarse = coarse_model
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_or_path, filename=None, coarse_filename=None, cache_dir=None, local_files_only=False):
-        filename = filename or "sk_model.pth"
-        coarse_filename = coarse_filename or "sk_model2.pth"
-
-        if os.path.isdir(pretrained_model_or_path):
-            model_path = os.path.join(pretrained_model_or_path, filename)
-            coarse_model_path = os.path.join(pretrained_model_or_path, coarse_filename)
-        else:
-            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir, local_files_only=local_files_only)
-            coarse_model_path = hf_hub_download(pretrained_model_or_path, coarse_filename, cache_dir=cache_dir, local_files_only=local_files_only)
-
-        model = Generator(3, 1, 3)
-        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
-        model.eval()
-
-        coarse_model = Generator(3, 1, 3)
-        coarse_model.load_state_dict(torch.load(coarse_model_path, map_location=torch.device('cpu')))
-        coarse_model.eval()
-
-        return cls(model, coarse_model)
-    
-    def to(self, device):
-        self.model.to(device)
-        self.model_coarse.to(device)
-        return self
-    
-    def __call__(self, input_image, coarse=False, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
-        if "return_pil" in kwargs:
-            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
-            output_type = "pil" if kwargs["return_pil"] else "np"
-        if type(output_type) is bool:
-            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
-            if output_type:
-                output_type = "pil"
-
-        device = next(iter(self.model.parameters())).device
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-
-        model = self.model_coarse if coarse else self.model
-        assert input_image.ndim == 3
-        image = input_image
-        with torch.no_grad():
-            image = torch.from_numpy(image).float().to(device)
-            image = image / 255.0
-            image = rearrange(image, 'h w c -> 1 c h w')
-            line = model(image)[0][0]
-
-            line = line.cpu().numpy()
-            line = (line * 255.0).clip(0, 255).astype(np.uint8)
-
-        detected_map = line
-
-        detected_map = HWC3(detected_map)
-
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-        detected_map = 255 - detected_map
-        
-        if output_type == "pil":
-            detected_map = Image.fromarray(detected_map)
-
-        return detected_map

controlnet_aux_local/lineart_anime/__init__.py

@@ -1,189 +0,0 @@
-import functools
-import os
-import warnings
-
-import cv2
-import numpy as np
-import torch
-import torch.nn as nn
-from einops import rearrange
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from ..util import HWC3, resize_image
-
-
-class UnetGenerator(nn.Module):
-    """Create a Unet-based generator"""
-
-    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
-        """Construct a Unet generator
-        Parameters:
-            input_nc (int)  -- the number of channels in input images
-            output_nc (int) -- the number of channels in output images
-            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
-                                image of size 128x128 will become of size 1x1 # at the bottleneck
-            ngf (int)       -- the number of filters in the last conv layer
-            norm_layer      -- normalization layer
-        We construct the U-Net from the innermost layer to the outermost layer.
-        It is a recursive process.
-        """
-        super(UnetGenerator, self).__init__()
-        # construct unet structure
-        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer
-        for _ in range(num_downs - 5):          # add intermediate layers with ngf * 8 filters
-            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
-        # gradually reduce the number of filters from ngf * 8 to ngf
-        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
-        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer
-
-    def forward(self, input):
-        """Standard forward"""
-        return self.model(input)
-
-
-class UnetSkipConnectionBlock(nn.Module):
-    """Defines the Unet submodule with skip connection.
-        X -------------------identity----------------------
-        |-- downsampling -- |submodule| -- upsampling --|
-    """
-
-    def __init__(self, outer_nc, inner_nc, input_nc=None,
-                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
-        """Construct a Unet submodule with skip connections.
-        Parameters:
-            outer_nc (int) -- the number of filters in the outer conv layer
-            inner_nc (int) -- the number of filters in the inner conv layer
-            input_nc (int) -- the number of channels in input images/features
-            submodule (UnetSkipConnectionBlock) -- previously defined submodules
-            outermost (bool)    -- if this module is the outermost module
-            innermost (bool)    -- if this module is the innermost module
-            norm_layer          -- normalization layer
-            use_dropout (bool)  -- if use dropout layers.
-        """
-        super(UnetSkipConnectionBlock, self).__init__()
-        self.outermost = outermost
-        if type(norm_layer) == functools.partial:
-            use_bias = norm_layer.func == nn.InstanceNorm2d
-        else:
-            use_bias = norm_layer == nn.InstanceNorm2d
-        if input_nc is None:
-            input_nc = outer_nc
-        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
-                             stride=2, padding=1, bias=use_bias)
-        downrelu = nn.LeakyReLU(0.2, True)
-        downnorm = norm_layer(inner_nc)
-        uprelu = nn.ReLU(True)
-        upnorm = norm_layer(outer_nc)
-
-        if outermost:
-            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1)
-            down = [downconv]
-            up = [uprelu, upconv, nn.Tanh()]
-            model = down + [submodule] + up
-        elif innermost:
-            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1, bias=use_bias)
-            down = [downrelu, downconv]
-            up = [uprelu, upconv, upnorm]
-            model = down + up
-        else:
-            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
-                                        kernel_size=4, stride=2,
-                                        padding=1, bias=use_bias)
-            down = [downrelu, downconv, downnorm]
-            up = [uprelu, upconv, upnorm]
-
-            if use_dropout:
-                model = down + [submodule] + up + [nn.Dropout(0.5)]
-            else:
-                model = down + [submodule] + up
-
-        self.model = nn.Sequential(*model)
-
-    def forward(self, x):
-        if self.outermost:
-            return self.model(x)
-        else:   # add skip connections
-            return torch.cat([x, self.model(x)], 1)
-
-
-class LineartAnimeDetector:
-    def __init__(self, model):
-        self.model = model
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None, local_files_only=False):
-        filename = filename or "netG.pth"
-
-        if os.path.isdir(pretrained_model_or_path):
-            model_path = os.path.join(pretrained_model_or_path, filename)
-        else:
-            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir, local_files_only=local_files_only)
-
-        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
-        net = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False)
-        ckpt = torch.load(model_path)
-        for key in list(ckpt.keys()):
-            if 'module.' in key:
-                ckpt[key.replace('module.', '')] = ckpt[key]
-                del ckpt[key]
-        net.load_state_dict(ckpt)
-        net.eval()
-
-        return cls(net)
-
-    def to(self, device):
-        self.model.to(device)
-        return self
-    
-    def __call__(self, input_image, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
-        if "return_pil" in kwargs:
-            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
-            output_type = "pil" if kwargs["return_pil"] else "np"
-        if type(output_type) is bool:
-            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
-            if output_type:
-                output_type = "pil"
-
-        device = next(iter(self.model.parameters())).device
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-
-        H, W, C = input_image.shape
-        Hn = 256 * int(np.ceil(float(H) / 256.0))
-        Wn = 256 * int(np.ceil(float(W) / 256.0))
-        img = cv2.resize(input_image, (Wn, Hn), interpolation=cv2.INTER_CUBIC)
-        with torch.no_grad():
-            image_feed = torch.from_numpy(img).float().to(device)
-            image_feed = image_feed / 127.5 - 1.0
-            image_feed = rearrange(image_feed, 'h w c -> 1 c h w')
-
-            line = self.model(image_feed)[0, 0] * 127.5 + 127.5
-            line = line.cpu().numpy()
-
-            line = cv2.resize(line, (W, H), interpolation=cv2.INTER_CUBIC)
-            line = line.clip(0, 255).astype(np.uint8)
-
-        detected_map = line
-
-        detected_map = HWC3(detected_map)
-
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-        detected_map = 255 - detected_map
-        
-        if output_type == "pil":
-            detected_map = Image.fromarray(detected_map)
-
-        return detected_map

controlnet_aux_local/mediapipe_face/__init__.py

@@ -1,53 +0,0 @@
-import warnings
-from typing import Union
-
-import cv2
-import numpy as np
-from PIL import Image
-
-from ..util import HWC3, resize_image
-from .mediapipe_face_common import generate_annotation
-
-
-class MediapipeFaceDetector:
-    def __call__(self,
-                 input_image: Union[np.ndarray, Image.Image] = None,
-                 max_faces: int = 1,
-                 min_confidence: float = 0.5,
-                 output_type: str = "pil",
-                 detect_resolution: int = 512,
-                 image_resolution: int = 512,
-                 **kwargs):
-
-        if "image" in kwargs:
-            warnings.warn("image is deprecated, please use `input_image=...` instead.", DeprecationWarning)
-            input_image = kwargs.pop("image")        
-        if input_image is None:
-            raise ValueError("input_image must be defined.")
-
-        if "return_pil" in kwargs:
-            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
-            output_type = "pil" if kwargs["return_pil"] else "np"
-        if type(output_type) is bool:
-            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
-            if output_type:
-                output_type = "pil"
-
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-
-        detected_map = generate_annotation(input_image, max_faces, min_confidence)
-        detected_map = HWC3(detected_map)
-
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-        
-        if output_type == "pil":
-            detected_map = Image.fromarray(detected_map)
-            
-        return detected_map

controlnet_aux_local/mediapipe_face/mediapipe_face_common.py

@@ -1,164 +0,0 @@
-from typing import Mapping
-import warnings
-
-try:
-    import mediapipe as mp
-except ImportError:
-    warnings.warn(
-        "The module 'mediapipe' is not installed. The package will have limited functionality. Please install it using the command: pip install 'mediapipe'"
-    )
-
-    mp = None
-
-import numpy
-
-if mp:
-    mp_drawing = mp.solutions.drawing_utils
-    mp_drawing_styles = mp.solutions.drawing_styles
-    mp_face_detection = mp.solutions.face_detection  # Only for counting faces.
-    mp_face_mesh = mp.solutions.face_mesh
-    mp_face_connections = mp.solutions.face_mesh_connections.FACEMESH_TESSELATION
-    mp_hand_connections = mp.solutions.hands_connections.HAND_CONNECTIONS
-    mp_body_connections = mp.solutions.pose_connections.POSE_CONNECTIONS
-
-    DrawingSpec = mp.solutions.drawing_styles.DrawingSpec
-    PoseLandmark = mp.solutions.drawing_styles.PoseLandmark
-
-    min_face_size_pixels: int = 64
-    f_thick = 2
-    f_rad = 1
-    right_iris_draw = DrawingSpec(color=(10, 200, 250), thickness=f_thick, circle_radius=f_rad)
-    right_eye_draw = DrawingSpec(color=(10, 200, 180), thickness=f_thick, circle_radius=f_rad)
-    right_eyebrow_draw = DrawingSpec(color=(10, 220, 180), thickness=f_thick, circle_radius=f_rad)
-    left_iris_draw = DrawingSpec(color=(250, 200, 10), thickness=f_thick, circle_radius=f_rad)
-    left_eye_draw = DrawingSpec(color=(180, 200, 10), thickness=f_thick, circle_radius=f_rad)
-    left_eyebrow_draw = DrawingSpec(color=(180, 220, 10), thickness=f_thick, circle_radius=f_rad)
-    mouth_draw = DrawingSpec(color=(10, 180, 10), thickness=f_thick, circle_radius=f_rad)
-    head_draw = DrawingSpec(color=(10, 200, 10), thickness=f_thick, circle_radius=f_rad)
-
-    # mp_face_mesh.FACEMESH_CONTOURS has all the items we care about.
-    face_connection_spec = {}
-    for edge in mp_face_mesh.FACEMESH_FACE_OVAL:
-        face_connection_spec[edge] = head_draw
-    for edge in mp_face_mesh.FACEMESH_LEFT_EYE:
-        face_connection_spec[edge] = left_eye_draw
-    for edge in mp_face_mesh.FACEMESH_LEFT_EYEBROW:
-        face_connection_spec[edge] = left_eyebrow_draw
-    # for edge in mp_face_mesh.FACEMESH_LEFT_IRIS:
-    #    face_connection_spec[edge] = left_iris_draw
-    for edge in mp_face_mesh.FACEMESH_RIGHT_EYE:
-        face_connection_spec[edge] = right_eye_draw
-    for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
-        face_connection_spec[edge] = right_eyebrow_draw
-    # for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:
-    #    face_connection_spec[edge] = right_iris_draw
-    for edge in mp_face_mesh.FACEMESH_LIPS:
-        face_connection_spec[edge] = mouth_draw
-    iris_landmark_spec = {468: right_iris_draw, 473: left_iris_draw}
-
-
-def draw_pupils(image, landmark_list, drawing_spec, halfwidth: int = 2):
-    """We have a custom function to draw the pupils because the mp.draw_landmarks method requires a parameter for all
-    landmarks.  Until our PR is merged into mediapipe, we need this separate method."""
-    if len(image.shape) != 3:
-        raise ValueError("Input image must be H,W,C.")
-    image_rows, image_cols, image_channels = image.shape
-    if image_channels != 3:  # BGR channels
-        raise ValueError('Input image must contain three channel bgr data.')
-    for idx, landmark in enumerate(landmark_list.landmark):
-        if (
-                (landmark.HasField('visibility') and landmark.visibility < 0.9) or
-                (landmark.HasField('presence') and landmark.presence < 0.5)
-        ):
-            continue
-        if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:
-            continue
-        image_x = int(image_cols*landmark.x)
-        image_y = int(image_rows*landmark.y)
-        draw_color = None
-        if isinstance(drawing_spec, Mapping):
-            if drawing_spec.get(idx) is None:
-                continue
-            else:
-                draw_color = drawing_spec[idx].color
-        elif isinstance(drawing_spec, DrawingSpec):
-            draw_color = drawing_spec.color
-        image[image_y-halfwidth:image_y+halfwidth, image_x-halfwidth:image_x+halfwidth, :] = draw_color
-
-
-def reverse_channels(image):
-    """Given a numpy array in RGB form, convert to BGR.  Will also convert from BGR to RGB."""
-    # im[:,:,::-1] is a neat hack to convert BGR to RGB by reversing the indexing order.
-    # im[:,:,::[2,1,0]] would also work but makes a copy of the data.
-    return image[:, :, ::-1]
-
-
-def generate_annotation(
-        img_rgb,
-        max_faces: int,
-        min_confidence: float
-):
-    """
-    Find up to 'max_faces' inside the provided input image.
-    If min_face_size_pixels is provided and nonzero it will be used to filter faces that occupy less than this many
-    pixels in the image.
-    """
-    with mp_face_mesh.FaceMesh(
-            static_image_mode=True,
-            max_num_faces=max_faces,
-            refine_landmarks=True,
-            min_detection_confidence=min_confidence,
-    ) as facemesh:
-        img_height, img_width, img_channels = img_rgb.shape
-        assert(img_channels == 3)
-
-        results = facemesh.process(img_rgb).multi_face_landmarks
-
-        if results is None:
-            print("No faces detected in controlnet image for Mediapipe face annotator.")
-            return numpy.zeros_like(img_rgb)
-
-        # Filter faces that are too small
-        filtered_landmarks = []
-        for lm in results:
-            landmarks = lm.landmark
-            face_rect = [
-                landmarks[0].x,
-                landmarks[0].y,
-                landmarks[0].x,
-                landmarks[0].y,
-            ]  # Left, up, right, down.
-            for i in range(len(landmarks)):
-                face_rect[0] = min(face_rect[0], landmarks[i].x)
-                face_rect[1] = min(face_rect[1], landmarks[i].y)
-                face_rect[2] = max(face_rect[2], landmarks[i].x)
-                face_rect[3] = max(face_rect[3], landmarks[i].y)
-            if min_face_size_pixels > 0:
-                face_width = abs(face_rect[2] - face_rect[0])
-                face_height = abs(face_rect[3] - face_rect[1])
-                face_width_pixels = face_width * img_width
-                face_height_pixels = face_height * img_height
-                face_size = min(face_width_pixels, face_height_pixels)
-                if face_size >= min_face_size_pixels:
-                    filtered_landmarks.append(lm)
-            else:
-                filtered_landmarks.append(lm)
-
-        # Annotations are drawn in BGR for some reason, but we don't need to flip a zero-filled image at the start.
-        empty = numpy.zeros_like(img_rgb)
-
-        # Draw detected faces:
-        for face_landmarks in filtered_landmarks:
-            mp_drawing.draw_landmarks(
-                empty,
-                face_landmarks,
-                connections=face_connection_spec.keys(),
-                landmark_drawing_spec=None,
-                connection_drawing_spec=face_connection_spec
-            )
-            draw_pupils(empty, face_landmarks, iris_landmark_spec, 2)
-
-        # Flip BGR back to RGB.
-        empty = reverse_channels(empty).copy()
-
-        return empty

controlnet_aux_local/midas/__init__.py

@@ -1,95 +0,0 @@
-import os
-
-import cv2
-import numpy as np
-import torch
-from einops import rearrange
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from ..util import HWC3, resize_image
-from .api import MiDaSInference
-
-
-class MidasDetector:
-    def __init__(self, model):
-        self.model = model
-        
-    @classmethod
-    def from_pretrained(cls, pretrained_model_or_path, model_type="dpt_hybrid", filename=None, cache_dir=None, local_files_only=False):
-        if pretrained_model_or_path == "lllyasviel/ControlNet":
-            filename = filename or "annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
-        else:
-            filename = filename or "dpt_hybrid-midas-501f0c75.pt"
-
-        if os.path.isdir(pretrained_model_or_path):
-            model_path = os.path.join(pretrained_model_or_path, filename)
-        else:
-            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir, local_files_only=local_files_only)
-
-        model = MiDaSInference(model_type=model_type, model_path=model_path)
-
-        return cls(model)
-        
-
-    def to(self, device):
-        self.model.to(device)
-        return self
-    
-    def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1, depth_and_normal=False, detect_resolution=512, image_resolution=512, output_type=None):
-        device = next(iter(self.model.parameters())).device
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-            output_type = output_type or "pil"
-        else:
-            output_type = output_type or "np"
-        
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-
-        assert input_image.ndim == 3
-        image_depth = input_image
-        with torch.no_grad():
-            image_depth = torch.from_numpy(image_depth).float()
-            image_depth = image_depth.to(device)
-            image_depth = image_depth / 127.5 - 1.0
-            image_depth = rearrange(image_depth, 'h w c -> 1 c h w')
-            depth = self.model(image_depth)[0]
-
-            depth_pt = depth.clone()
-            depth_pt -= torch.min(depth_pt)
-            depth_pt /= torch.max(depth_pt)
-            depth_pt = depth_pt.cpu().numpy()
-            depth_image = (depth_pt * 255.0).clip(0, 255).astype(np.uint8)
-
-            if depth_and_normal:
-                depth_np = depth.cpu().numpy()
-                x = cv2.Sobel(depth_np, cv2.CV_32F, 1, 0, ksize=3)
-                y = cv2.Sobel(depth_np, cv2.CV_32F, 0, 1, ksize=3)
-                z = np.ones_like(x) * a
-                x[depth_pt < bg_th] = 0
-                y[depth_pt < bg_th] = 0
-                normal = np.stack([x, y, z], axis=2)
-                normal /= np.sum(normal ** 2.0, axis=2, keepdims=True) ** 0.5
-                normal_image = (normal * 127.5 + 127.5).clip(0, 255).astype(np.uint8)[:, :, ::-1]
-        
-        depth_image = HWC3(depth_image)
-        if depth_and_normal:
-            normal_image = HWC3(normal_image)
-
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        depth_image = cv2.resize(depth_image, (W, H), interpolation=cv2.INTER_LINEAR)
-        if depth_and_normal:
-            normal_image = cv2.resize(normal_image, (W, H), interpolation=cv2.INTER_LINEAR)
-        
-        if output_type == "pil":
-            depth_image = Image.fromarray(depth_image)
-            if depth_and_normal:
-                normal_image = Image.fromarray(normal_image)
-        
-        if depth_and_normal:
-            return depth_image, normal_image
-        else:
-            return depth_image

controlnet_aux_local/midas/api.py

@@ -1,169 +0,0 @@
-# based on https://github.com/isl-org/MiDaS
-
-import cv2
-import os
-import torch
-import torch.nn as nn
-from torchvision.transforms import Compose
-
-from .midas.dpt_depth import DPTDepthModel
-from .midas.midas_net import MidasNet
-from .midas.midas_net_custom import MidasNet_small
-from .midas.transforms import Resize, NormalizeImage, PrepareForNet
-from ..util import annotator_ckpts_path
-
-
-ISL_PATHS = {
-    "dpt_large": os.path.join(annotator_ckpts_path, "dpt_large-midas-2f21e586.pt"),
-    "dpt_hybrid": os.path.join(annotator_ckpts_path, "dpt_hybrid-midas-501f0c75.pt"),
-    "midas_v21": "",
-    "midas_v21_small": "",
-}
-
-remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
-
-
-def disabled_train(self, mode=True):
-    """Overwrite model.train with this function to make sure train/eval mode
-    does not change anymore."""
-    return self
-
-
-def load_midas_transform(model_type):
-    # https://github.com/isl-org/MiDaS/blob/master/run.py
-    # load transform only
-    if model_type == "dpt_large":  # DPT-Large
-        net_w, net_h = 384, 384
-        resize_mode = "minimal"
-        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-    elif model_type == "dpt_hybrid":  # DPT-Hybrid
-        net_w, net_h = 384, 384
-        resize_mode = "minimal"
-        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-    elif model_type == "midas_v21":
-        net_w, net_h = 384, 384
-        resize_mode = "upper_bound"
-        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-
-    elif model_type == "midas_v21_small":
-        net_w, net_h = 256, 256
-        resize_mode = "upper_bound"
-        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-
-    else:
-        assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
-
-    transform = Compose(
-        [
-            Resize(
-                net_w,
-                net_h,
-                resize_target=None,
-                keep_aspect_ratio=True,
-                ensure_multiple_of=32,
-                resize_method=resize_mode,
-                image_interpolation_method=cv2.INTER_CUBIC,
-            ),
-            normalization,
-            PrepareForNet(),
-        ]
-    )
-
-    return transform
-
-
-def load_model(model_type, model_path=None):
-    # https://github.com/isl-org/MiDaS/blob/master/run.py
-    # load network
-    model_path = model_path or ISL_PATHS[model_type]
-    if model_type == "dpt_large":  # DPT-Large
-        model = DPTDepthModel(
-            path=model_path,
-            backbone="vitl16_384",
-            non_negative=True,
-        )
-        net_w, net_h = 384, 384
-        resize_mode = "minimal"
-        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-    elif model_type == "dpt_hybrid":  # DPT-Hybrid
-        if not os.path.exists(model_path):
-            from basicsr.utils.download_util import load_file_from_url
-            load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path)
-
-        model = DPTDepthModel(
-            path=model_path,
-            backbone="vitb_rn50_384",
-            non_negative=True,
-        )
-        net_w, net_h = 384, 384
-        resize_mode = "minimal"
-        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-    elif model_type == "midas_v21":
-        model = MidasNet(model_path, non_negative=True)
-        net_w, net_h = 384, 384
-        resize_mode = "upper_bound"
-        normalization = NormalizeImage(
-            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
-        )
-
-    elif model_type == "midas_v21_small":
-        model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True,
-                               non_negative=True, blocks={'expand': True})
-        net_w, net_h = 256, 256
-        resize_mode = "upper_bound"
-        normalization = NormalizeImage(
-            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
-        )
-
-    else:
-        print(f"model_type '{model_type}' not implemented, use: --model_type large")
-        assert False
-
-    transform = Compose(
-        [
-            Resize(
-                net_w,
-                net_h,
-                resize_target=None,
-                keep_aspect_ratio=True,
-                ensure_multiple_of=32,
-                resize_method=resize_mode,
-                image_interpolation_method=cv2.INTER_CUBIC,
-            ),
-            normalization,
-            PrepareForNet(),
-        ]
-    )
-
-    return model.eval(), transform
-
-
-class MiDaSInference(nn.Module):
-    MODEL_TYPES_TORCH_HUB = [
-        "DPT_Large",
-        "DPT_Hybrid",
-        "MiDaS_small"
-    ]
-    MODEL_TYPES_ISL = [
-        "dpt_large",
-        "dpt_hybrid",
-        "midas_v21",
-        "midas_v21_small",
-    ]
-
-    def __init__(self, model_type, model_path):
-        super().__init__()
-        assert (model_type in self.MODEL_TYPES_ISL)
-        model, _ = load_model(model_type, model_path)
-        self.model = model
-        self.model.train = disabled_train
-
-    def forward(self, x):
-        with torch.no_grad():
-            prediction = self.model(x)
-        return prediction
-

controlnet_aux_local/midas/midas/base_model.py

@@ -1,16 +0,0 @@
-import torch
-
-
-class BaseModel(torch.nn.Module):
-    def load(self, path):
-        """Load model from file.
-
-        Args:
-            path (str): file path
-        """
-        parameters = torch.load(path, map_location=torch.device('cpu'))
-
-        if "optimizer" in parameters:
-            parameters = parameters["model"]
-
-        self.load_state_dict(parameters)

controlnet_aux_local/midas/midas/blocks.py

@@ -1,342 +0,0 @@
-import torch
-import torch.nn as nn
-
-from .vit import (
-    _make_pretrained_vitb_rn50_384,
-    _make_pretrained_vitl16_384,
-    _make_pretrained_vitb16_384,
-    forward_vit,
-)
-
-def _make_encoder(backbone, features, use_pretrained, groups=1, expand=False, exportable=True, hooks=None, use_vit_only=False, use_readout="ignore",):
-    if backbone == "vitl16_384":
-        pretrained = _make_pretrained_vitl16_384(
-            use_pretrained, hooks=hooks, use_readout=use_readout
-        )
-        scratch = _make_scratch(
-            [256, 512, 1024, 1024], features, groups=groups, expand=expand
-        )  # ViT-L/16 - 85.0% Top1 (backbone)
-    elif backbone == "vitb_rn50_384":
-        pretrained = _make_pretrained_vitb_rn50_384(
-            use_pretrained,
-            hooks=hooks,
-            use_vit_only=use_vit_only,
-            use_readout=use_readout,
-        )
-        scratch = _make_scratch(
-            [256, 512, 768, 768], features, groups=groups, expand=expand
-        )  # ViT-H/16 - 85.0% Top1 (backbone)
-    elif backbone == "vitb16_384":
-        pretrained = _make_pretrained_vitb16_384(
-            use_pretrained, hooks=hooks, use_readout=use_readout
-        )
-        scratch = _make_scratch(
-            [96, 192, 384, 768], features, groups=groups, expand=expand
-        )  # ViT-B/16 - 84.6% Top1 (backbone)
-    elif backbone == "resnext101_wsl":
-        pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
-        scratch = _make_scratch([256, 512, 1024, 2048], features, groups=groups, expand=expand)     # efficientnet_lite3  
-    elif backbone == "efficientnet_lite3":
-        pretrained = _make_pretrained_efficientnet_lite3(use_pretrained, exportable=exportable)
-        scratch = _make_scratch([32, 48, 136, 384], features, groups=groups, expand=expand)  # efficientnet_lite3     
-    else:
-        print(f"Backbone '{backbone}' not implemented")
-        assert False
-        
-    return pretrained, scratch
-
-
-def _make_scratch(in_shape, out_shape, groups=1, expand=False):
-    scratch = nn.Module()
-
-    out_shape1 = out_shape
-    out_shape2 = out_shape
-    out_shape3 = out_shape
-    out_shape4 = out_shape
-    if expand==True:
-        out_shape1 = out_shape
-        out_shape2 = out_shape*2
-        out_shape3 = out_shape*4
-        out_shape4 = out_shape*8
-
-    scratch.layer1_rn = nn.Conv2d(
-        in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
-    )
-    scratch.layer2_rn = nn.Conv2d(
-        in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
-    )
-    scratch.layer3_rn = nn.Conv2d(
-        in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
-    )
-    scratch.layer4_rn = nn.Conv2d(
-        in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
-    )
-
-    return scratch
-
-
-def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
-    efficientnet = torch.hub.load(
-        "rwightman/gen-efficientnet-pytorch",
-        "tf_efficientnet_lite3",
-        pretrained=use_pretrained,
-        exportable=exportable
-    )
-    return _make_efficientnet_backbone(efficientnet)
-
-
-def _make_efficientnet_backbone(effnet):
-    pretrained = nn.Module()
-
-    pretrained.layer1 = nn.Sequential(
-        effnet.conv_stem, effnet.bn1, effnet.act1, *effnet.blocks[0:2]
-    )
-    pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
-    pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
-    pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])
-
-    return pretrained
-    
-
-def _make_resnet_backbone(resnet):
-    pretrained = nn.Module()
-    pretrained.layer1 = nn.Sequential(
-        resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1
-    )
-
-    pretrained.layer2 = resnet.layer2
-    pretrained.layer3 = resnet.layer3
-    pretrained.layer4 = resnet.layer4
-
-    return pretrained
-
-
-def _make_pretrained_resnext101_wsl(use_pretrained):
-    resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
-    return _make_resnet_backbone(resnet)
-
-
-
-class Interpolate(nn.Module):
-    """Interpolation module.
-    """
-
-    def __init__(self, scale_factor, mode, align_corners=False):
-        """Init.
-
-        Args:
-            scale_factor (float): scaling
-            mode (str): interpolation mode
-        """
-        super(Interpolate, self).__init__()
-
-        self.interp = nn.functional.interpolate
-        self.scale_factor = scale_factor
-        self.mode = mode
-        self.align_corners = align_corners
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input
-
-        Returns:
-            tensor: interpolated data
-        """
-
-        x = self.interp(
-            x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners
-        )
-
-        return x
-
-
-class ResidualConvUnit(nn.Module):
-    """Residual convolution module.
-    """
-
-    def __init__(self, features):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super().__init__()
-
-        self.conv1 = nn.Conv2d(
-            features, features, kernel_size=3, stride=1, padding=1, bias=True
-        )
-
-        self.conv2 = nn.Conv2d(
-            features, features, kernel_size=3, stride=1, padding=1, bias=True
-        )
-
-        self.relu = nn.ReLU(inplace=True)
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input
-
-        Returns:
-            tensor: output
-        """
-        out = self.relu(x)
-        out = self.conv1(out)
-        out = self.relu(out)
-        out = self.conv2(out)
-
-        return out + x
-
-
-class FeatureFusionBlock(nn.Module):
-    """Feature fusion block.
-    """
-
-    def __init__(self, features):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super(FeatureFusionBlock, self).__init__()
-
-        self.resConfUnit1 = ResidualConvUnit(features)
-        self.resConfUnit2 = ResidualConvUnit(features)
-
-    def forward(self, *xs):
-        """Forward pass.
-
-        Returns:
-            tensor: output
-        """
-        output = xs[0]
-
-        if len(xs) == 2:
-            output += self.resConfUnit1(xs[1])
-
-        output = self.resConfUnit2(output)
-
-        output = nn.functional.interpolate(
-            output, scale_factor=2, mode="bilinear", align_corners=True
-        )
-
-        return output
-
-
-
-
-class ResidualConvUnit_custom(nn.Module):
-    """Residual convolution module.
-    """
-
-    def __init__(self, features, activation, bn):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super().__init__()
-
-        self.bn = bn
-
-        self.groups=1
-
-        self.conv1 = nn.Conv2d(
-            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
-        )
-        
-        self.conv2 = nn.Conv2d(
-            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
-        )
-
-        if self.bn==True:
-            self.bn1 = nn.BatchNorm2d(features)
-            self.bn2 = nn.BatchNorm2d(features)
-
-        self.activation = activation
-
-        self.skip_add = nn.quantized.FloatFunctional()
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input
-
-        Returns:
-            tensor: output
-        """
-        
-        out = self.activation(x)
-        out = self.conv1(out)
-        if self.bn==True:
-            out = self.bn1(out)
-       
-        out = self.activation(out)
-        out = self.conv2(out)
-        if self.bn==True:
-            out = self.bn2(out)
-
-        if self.groups > 1:
-            out = self.conv_merge(out)
-
-        return self.skip_add.add(out, x)
-
-        # return out + x
-
-
-class FeatureFusionBlock_custom(nn.Module):
-    """Feature fusion block.
-    """
-
-    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super(FeatureFusionBlock_custom, self).__init__()
-
-        self.deconv = deconv
-        self.align_corners = align_corners
-
-        self.groups=1
-
-        self.expand = expand
-        out_features = features
-        if self.expand==True:
-            out_features = features//2
-        
-        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
-
-        self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
-        self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
-        
-        self.skip_add = nn.quantized.FloatFunctional()
-
-    def forward(self, *xs):
-        """Forward pass.
-
-        Returns:
-            tensor: output
-        """
-        output = xs[0]
-
-        if len(xs) == 2:
-            res = self.resConfUnit1(xs[1])
-            output = self.skip_add.add(output, res)
-            # output += res
-
-        output = self.resConfUnit2(output)
-
-        output = nn.functional.interpolate(
-            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
-        )
-
-        output = self.out_conv(output)
-
-        return output
-

controlnet_aux_local/midas/midas/dpt_depth.py

@@ -1,109 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .base_model import BaseModel
-from .blocks import (
-    FeatureFusionBlock,
-    FeatureFusionBlock_custom,
-    Interpolate,
-    _make_encoder,
-    forward_vit,
-)
-
-
-def _make_fusion_block(features, use_bn):
-    return FeatureFusionBlock_custom(
-        features,
-        nn.ReLU(False),
-        deconv=False,
-        bn=use_bn,
-        expand=False,
-        align_corners=True,
-    )
-
-
-class DPT(BaseModel):
-    def __init__(
-        self,
-        head,
-        features=256,
-        backbone="vitb_rn50_384",
-        readout="project",
-        channels_last=False,
-        use_bn=False,
-    ):
-
-        super(DPT, self).__init__()
-
-        self.channels_last = channels_last
-
-        hooks = {
-            "vitb_rn50_384": [0, 1, 8, 11],
-            "vitb16_384": [2, 5, 8, 11],
-            "vitl16_384": [5, 11, 17, 23],
-        }
-
-        # Instantiate backbone and reassemble blocks
-        self.pretrained, self.scratch = _make_encoder(
-            backbone,
-            features,
-            False, # Set to true of you want to train from scratch, uses ImageNet weights
-            groups=1,
-            expand=False,
-            exportable=False,
-            hooks=hooks[backbone],
-            use_readout=readout,
-        )
-
-        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
-        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
-        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
-        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
-
-        self.scratch.output_conv = head
-
-
-    def forward(self, x):
-        if self.channels_last == True:
-            x.contiguous(memory_format=torch.channels_last)
-
-        layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)
-
-        layer_1_rn = self.scratch.layer1_rn(layer_1)
-        layer_2_rn = self.scratch.layer2_rn(layer_2)
-        layer_3_rn = self.scratch.layer3_rn(layer_3)
-        layer_4_rn = self.scratch.layer4_rn(layer_4)
-
-        path_4 = self.scratch.refinenet4(layer_4_rn)
-        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
-        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
-        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
-
-        out = self.scratch.output_conv(path_1)
-
-        return out
-
-
-class DPTDepthModel(DPT):
-    def __init__(self, path=None, non_negative=True, **kwargs):
-        features = kwargs["features"] if "features" in kwargs else 256
-
-        head = nn.Sequential(
-            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
-            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
-            nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(True),
-            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
-            nn.ReLU(True) if non_negative else nn.Identity(),
-            nn.Identity(),
-        )
-
-        super().__init__(head, **kwargs)
-
-        if path is not None:
-           self.load(path)
-
-    def forward(self, x):
-        return super().forward(x).squeeze(dim=1)
-

controlnet_aux_local/midas/midas/midas_net.py

@@ -1,76 +0,0 @@
-"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
-This file contains code that is adapted from
-https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
-"""
-import torch
-import torch.nn as nn
-
-from .base_model import BaseModel
-from .blocks import FeatureFusionBlock, Interpolate, _make_encoder
-
-
-class MidasNet(BaseModel):
-    """Network for monocular depth estimation.
-    """
-
-    def __init__(self, path=None, features=256, non_negative=True):
-        """Init.
-
-        Args:
-            path (str, optional): Path to saved model. Defaults to None.
-            features (int, optional): Number of features. Defaults to 256.
-            backbone (str, optional): Backbone network for encoder. Defaults to resnet50
-        """
-        print("Loading weights: ", path)
-
-        super(MidasNet, self).__init__()
-
-        use_pretrained = False if path is None else True
-
-        self.pretrained, self.scratch = _make_encoder(backbone="resnext101_wsl", features=features, use_pretrained=use_pretrained)
-
-        self.scratch.refinenet4 = FeatureFusionBlock(features)
-        self.scratch.refinenet3 = FeatureFusionBlock(features)
-        self.scratch.refinenet2 = FeatureFusionBlock(features)
-        self.scratch.refinenet1 = FeatureFusionBlock(features)
-
-        self.scratch.output_conv = nn.Sequential(
-            nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1),
-            Interpolate(scale_factor=2, mode="bilinear"),
-            nn.Conv2d(128, 32, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(True),
-            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
-            nn.ReLU(True) if non_negative else nn.Identity(),
-        )
-
-        if path:
-            self.load(path)
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input data (image)
-
-        Returns:
-            tensor: depth
-        """
-
-        layer_1 = self.pretrained.layer1(x)
-        layer_2 = self.pretrained.layer2(layer_1)
-        layer_3 = self.pretrained.layer3(layer_2)
-        layer_4 = self.pretrained.layer4(layer_3)
-
-        layer_1_rn = self.scratch.layer1_rn(layer_1)
-        layer_2_rn = self.scratch.layer2_rn(layer_2)
-        layer_3_rn = self.scratch.layer3_rn(layer_3)
-        layer_4_rn = self.scratch.layer4_rn(layer_4)
-
-        path_4 = self.scratch.refinenet4(layer_4_rn)
-        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
-        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
-        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
-
-        out = self.scratch.output_conv(path_1)
-
-        return torch.squeeze(out, dim=1)

controlnet_aux_local/midas/midas/midas_net_custom.py

@@ -1,128 +0,0 @@
-"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
-This file contains code that is adapted from
-https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
-"""
-import torch
-import torch.nn as nn
-
-from .base_model import BaseModel
-from .blocks import FeatureFusionBlock, FeatureFusionBlock_custom, Interpolate, _make_encoder
-
-
-class MidasNet_small(BaseModel):
-    """Network for monocular depth estimation.
-    """
-
-    def __init__(self, path=None, features=64, backbone="efficientnet_lite3", non_negative=True, exportable=True, channels_last=False, align_corners=True,
-        blocks={'expand': True}):
-        """Init.
-
-        Args:
-            path (str, optional): Path to saved model. Defaults to None.
-            features (int, optional): Number of features. Defaults to 256.
-            backbone (str, optional): Backbone network for encoder. Defaults to resnet50
-        """
-        print("Loading weights: ", path)
-
-        super(MidasNet_small, self).__init__()
-
-        use_pretrained = False if path else True
-                
-        self.channels_last = channels_last
-        self.blocks = blocks
-        self.backbone = backbone
-
-        self.groups = 1
-
-        features1=features
-        features2=features
-        features3=features
-        features4=features
-        self.expand = False
-        if "expand" in self.blocks and self.blocks['expand'] == True:
-            self.expand = True
-            features1=features
-            features2=features*2
-            features3=features*4
-            features4=features*8
-
-        self.pretrained, self.scratch = _make_encoder(self.backbone, features, use_pretrained, groups=self.groups, expand=self.expand, exportable=exportable)
-  
-        self.scratch.activation = nn.ReLU(False)    
-
-        self.scratch.refinenet4 = FeatureFusionBlock_custom(features4, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
-        self.scratch.refinenet3 = FeatureFusionBlock_custom(features3, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
-        self.scratch.refinenet2 = FeatureFusionBlock_custom(features2, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
-        self.scratch.refinenet1 = FeatureFusionBlock_custom(features1, self.scratch.activation, deconv=False, bn=False, align_corners=align_corners)
-
-        
-        self.scratch.output_conv = nn.Sequential(
-            nn.Conv2d(features, features//2, kernel_size=3, stride=1, padding=1, groups=self.groups),
-            Interpolate(scale_factor=2, mode="bilinear"),
-            nn.Conv2d(features//2, 32, kernel_size=3, stride=1, padding=1),
-            self.scratch.activation,
-            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
-            nn.ReLU(True) if non_negative else nn.Identity(),
-            nn.Identity(),
-        )
-        
-        if path:
-            self.load(path)
-
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input data (image)
-
-        Returns:
-            tensor: depth
-        """
-        if self.channels_last==True:
-            print("self.channels_last = ", self.channels_last)
-            x.contiguous(memory_format=torch.channels_last)
-
-
-        layer_1 = self.pretrained.layer1(x)
-        layer_2 = self.pretrained.layer2(layer_1)
-        layer_3 = self.pretrained.layer3(layer_2)
-        layer_4 = self.pretrained.layer4(layer_3)
-        
-        layer_1_rn = self.scratch.layer1_rn(layer_1)
-        layer_2_rn = self.scratch.layer2_rn(layer_2)
-        layer_3_rn = self.scratch.layer3_rn(layer_3)
-        layer_4_rn = self.scratch.layer4_rn(layer_4)
-
-
-        path_4 = self.scratch.refinenet4(layer_4_rn)
-        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
-        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
-        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
-        
-        out = self.scratch.output_conv(path_1)
-
-        return torch.squeeze(out, dim=1)
-
-
-
-def fuse_model(m):
-    prev_previous_type = nn.Identity()
-    prev_previous_name = ''
-    previous_type = nn.Identity()
-    previous_name = ''
-    for name, module in m.named_modules():
-        if prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d and type(module) == nn.ReLU:
-            # print("FUSED ", prev_previous_name, previous_name, name)
-            torch.quantization.fuse_modules(m, [prev_previous_name, previous_name, name], inplace=True)
-        elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d:
-            # print("FUSED ", prev_previous_name, previous_name)
-            torch.quantization.fuse_modules(m, [prev_previous_name, previous_name], inplace=True)
-        # elif previous_type == nn.Conv2d and type(module) == nn.ReLU:
-        #    print("FUSED ", previous_name, name)
-        #    torch.quantization.fuse_modules(m, [previous_name, name], inplace=True)
-
-        prev_previous_type = previous_type
-        prev_previous_name = previous_name
-        previous_type = type(module)
-        previous_name = name

controlnet_aux_local/midas/midas/transforms.py

@@ -1,234 +0,0 @@
-import numpy as np
-import cv2
-import math
-
-
-def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
-    """Rezise the sample to ensure the given size. Keeps aspect ratio.
-
-    Args:
-        sample (dict): sample
-        size (tuple): image size
-
-    Returns:
-        tuple: new size
-    """
-    shape = list(sample["disparity"].shape)
-
-    if shape[0] >= size[0] and shape[1] >= size[1]:
-        return sample
-
-    scale = [0, 0]
-    scale[0] = size[0] / shape[0]
-    scale[1] = size[1] / shape[1]
-
-    scale = max(scale)
-
-    shape[0] = math.ceil(scale * shape[0])
-    shape[1] = math.ceil(scale * shape[1])
-
-    # resize
-    sample["image"] = cv2.resize(
-        sample["image"], tuple(shape[::-1]), interpolation=image_interpolation_method
-    )
-
-    sample["disparity"] = cv2.resize(
-        sample["disparity"], tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST
-    )
-    sample["mask"] = cv2.resize(
-        sample["mask"].astype(np.float32),
-        tuple(shape[::-1]),
-        interpolation=cv2.INTER_NEAREST,
-    )
-    sample["mask"] = sample["mask"].astype(bool)
-
-    return tuple(shape)
-
-
-class Resize(object):
-    """Resize sample to given size (width, height).
-    """
-
-    def __init__(
-        self,
-        width,
-        height,
-        resize_target=True,
-        keep_aspect_ratio=False,
-        ensure_multiple_of=1,
-        resize_method="lower_bound",
-        image_interpolation_method=cv2.INTER_AREA,
-    ):
-        """Init.
-
-        Args:
-            width (int): desired output width
-            height (int): desired output height
-            resize_target (bool, optional):
-                True: Resize the full sample (image, mask, target).
-                False: Resize image only.
-                Defaults to True.
-            keep_aspect_ratio (bool, optional):
-                True: Keep the aspect ratio of the input sample.
-                Output sample might not have the given width and height, and
-                resize behaviour depends on the parameter 'resize_method'.
-                Defaults to False.
-            ensure_multiple_of (int, optional):
-                Output width and height is constrained to be multiple of this parameter.
-                Defaults to 1.
-            resize_method (str, optional):
-                "lower_bound": Output will be at least as large as the given size.
-                "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
-                "minimal": Scale as least as possible.  (Output size might be smaller than given size.)
-                Defaults to "lower_bound".
-        """
-        self.__width = width
-        self.__height = height
-
-        self.__resize_target = resize_target
-        self.__keep_aspect_ratio = keep_aspect_ratio
-        self.__multiple_of = ensure_multiple_of
-        self.__resize_method = resize_method
-        self.__image_interpolation_method = image_interpolation_method
-
-    def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
-        y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
-        if max_val is not None and y > max_val:
-            y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
-        if y < min_val:
-            y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
-        return y
-
-    def get_size(self, width, height):
-        # determine new height and width
-        scale_height = self.__height / height
-        scale_width = self.__width / width
-
-        if self.__keep_aspect_ratio:
-            if self.__resize_method == "lower_bound":
-                # scale such that output size is lower bound
-                if scale_width > scale_height:
-                    # fit width
-                    scale_height = scale_width
-                else:
-                    # fit height
-                    scale_width = scale_height
-            elif self.__resize_method == "upper_bound":
-                # scale such that output size is upper bound
-                if scale_width < scale_height:
-                    # fit width
-                    scale_height = scale_width
-                else:
-                    # fit height
-                    scale_width = scale_height
-            elif self.__resize_method == "minimal":
-                # scale as least as possbile
-                if abs(1 - scale_width) < abs(1 - scale_height):
-                    # fit width
-                    scale_height = scale_width
-                else:
-                    # fit height
-                    scale_width = scale_height
-            else:
-                raise ValueError(
-                    f"resize_method {self.__resize_method} not implemented"
-                )
-
-        if self.__resize_method == "lower_bound":
-            new_height = self.constrain_to_multiple_of(
-                scale_height * height, min_val=self.__height
-            )
-            new_width = self.constrain_to_multiple_of(
-                scale_width * width, min_val=self.__width
-            )
-        elif self.__resize_method == "upper_bound":
-            new_height = self.constrain_to_multiple_of(
-                scale_height * height, max_val=self.__height
-            )
-            new_width = self.constrain_to_multiple_of(
-                scale_width * width, max_val=self.__width
-            )
-        elif self.__resize_method == "minimal":
-            new_height = self.constrain_to_multiple_of(scale_height * height)
-            new_width = self.constrain_to_multiple_of(scale_width * width)
-        else:
-            raise ValueError(f"resize_method {self.__resize_method} not implemented")
-
-        return (new_width, new_height)
-
-    def __call__(self, sample):
-        width, height = self.get_size(
-            sample["image"].shape[1], sample["image"].shape[0]
-        )
-
-        # resize sample
-        sample["image"] = cv2.resize(
-            sample["image"],
-            (width, height),
-            interpolation=self.__image_interpolation_method,
-        )
-
-        if self.__resize_target:
-            if "disparity" in sample:
-                sample["disparity"] = cv2.resize(
-                    sample["disparity"],
-                    (width, height),
-                    interpolation=cv2.INTER_NEAREST,
-                )
-
-            if "depth" in sample:
-                sample["depth"] = cv2.resize(
-                    sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST
-                )
-
-            sample["mask"] = cv2.resize(
-                sample["mask"].astype(np.float32),
-                (width, height),
-                interpolation=cv2.INTER_NEAREST,
-            )
-            sample["mask"] = sample["mask"].astype(bool)
-
-        return sample
-
-
-class NormalizeImage(object):
-    """Normlize image by given mean and std.
-    """
-
-    def __init__(self, mean, std):
-        self.__mean = mean
-        self.__std = std
-
-    def __call__(self, sample):
-        sample["image"] = (sample["image"] - self.__mean) / self.__std
-
-        return sample
-
-
-class PrepareForNet(object):
-    """Prepare sample for usage as network input.
-    """
-
-    def __init__(self):
-        pass
-
-    def __call__(self, sample):
-        image = np.transpose(sample["image"], (2, 0, 1))
-        sample["image"] = np.ascontiguousarray(image).astype(np.float32)
-
-        if "mask" in sample:
-            sample["mask"] = sample["mask"].astype(np.float32)
-            sample["mask"] = np.ascontiguousarray(sample["mask"])
-
-        if "disparity" in sample:
-            disparity = sample["disparity"].astype(np.float32)
-            sample["disparity"] = np.ascontiguousarray(disparity)
-
-        if "depth" in sample:
-            depth = sample["depth"].astype(np.float32)
-            sample["depth"] = np.ascontiguousarray(depth)
-
-        return sample

controlnet_aux_local/midas/midas/vit.py

@@ -1,491 +0,0 @@
-import torch
-import torch.nn as nn
-import timm
-import types
-import math
-import torch.nn.functional as F
-
-
-class Slice(nn.Module):
-    def __init__(self, start_index=1):
-        super(Slice, self).__init__()
-        self.start_index = start_index
-
-    def forward(self, x):
-        return x[:, self.start_index :]
-
-
-class AddReadout(nn.Module):
-    def __init__(self, start_index=1):
-        super(AddReadout, self).__init__()
-        self.start_index = start_index
-
-    def forward(self, x):
-        if self.start_index == 2:
-            readout = (x[:, 0] + x[:, 1]) / 2
-        else:
-            readout = x[:, 0]
-        return x[:, self.start_index :] + readout.unsqueeze(1)
-
-
-class ProjectReadout(nn.Module):
-    def __init__(self, in_features, start_index=1):
-        super(ProjectReadout, self).__init__()
-        self.start_index = start_index
-
-        self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), nn.GELU())
-
-    def forward(self, x):
-        readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index :])
-        features = torch.cat((x[:, self.start_index :], readout), -1)
-
-        return self.project(features)
-
-
-class Transpose(nn.Module):
-    def __init__(self, dim0, dim1):
-        super(Transpose, self).__init__()
-        self.dim0 = dim0
-        self.dim1 = dim1
-
-    def forward(self, x):
-        x = x.transpose(self.dim0, self.dim1)
-        return x
-
-
-def forward_vit(pretrained, x):
-    b, c, h, w = x.shape
-
-    glob = pretrained.model.forward_flex(x)
-
-    layer_1 = pretrained.activations["1"]
-    layer_2 = pretrained.activations["2"]
-    layer_3 = pretrained.activations["3"]
-    layer_4 = pretrained.activations["4"]
-
-    layer_1 = pretrained.act_postprocess1[0:2](layer_1)
-    layer_2 = pretrained.act_postprocess2[0:2](layer_2)
-    layer_3 = pretrained.act_postprocess3[0:2](layer_3)
-    layer_4 = pretrained.act_postprocess4[0:2](layer_4)
-
-    unflatten = nn.Sequential(
-        nn.Unflatten(
-            2,
-            torch.Size(
-                [
-                    h // pretrained.model.patch_size[1],
-                    w // pretrained.model.patch_size[0],
-                ]
-            ),
-        )
-    )
-
-    if layer_1.ndim == 3:
-        layer_1 = unflatten(layer_1)
-    if layer_2.ndim == 3:
-        layer_2 = unflatten(layer_2)
-    if layer_3.ndim == 3:
-        layer_3 = unflatten(layer_3)
-    if layer_4.ndim == 3:
-        layer_4 = unflatten(layer_4)
-
-    layer_1 = pretrained.act_postprocess1[3 : len(pretrained.act_postprocess1)](layer_1)
-    layer_2 = pretrained.act_postprocess2[3 : len(pretrained.act_postprocess2)](layer_2)
-    layer_3 = pretrained.act_postprocess3[3 : len(pretrained.act_postprocess3)](layer_3)
-    layer_4 = pretrained.act_postprocess4[3 : len(pretrained.act_postprocess4)](layer_4)
-
-    return layer_1, layer_2, layer_3, layer_4
-
-
-def _resize_pos_embed(self, posemb, gs_h, gs_w):
-    posemb_tok, posemb_grid = (
-        posemb[:, : self.start_index],
-        posemb[0, self.start_index :],
-    )
-
-    gs_old = int(math.sqrt(len(posemb_grid)))
-
-    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
-    posemb_grid = F.interpolate(posemb_grid, size=(gs_h, gs_w), mode="bilinear")
-    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1)
-
-    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
-
-    return posemb
-
-
-def forward_flex(self, x):
-    b, c, h, w = x.shape
-
-    pos_embed = self._resize_pos_embed(
-        self.pos_embed, h // self.patch_size[1], w // self.patch_size[0]
-    )
-
-    B = x.shape[0]
-
-    if hasattr(self.patch_embed, "backbone"):
-        x = self.patch_embed.backbone(x)
-        if isinstance(x, (list, tuple)):
-            x = x[-1]  # last feature if backbone outputs list/tuple of features
-
-    x = self.patch_embed.proj(x).flatten(2).transpose(1, 2)
-
-    if getattr(self, "dist_token", None) is not None:
-        cls_tokens = self.cls_token.expand(
-            B, -1, -1
-        )  # stole cls_tokens impl from Phil Wang, thanks
-        dist_token = self.dist_token.expand(B, -1, -1)
-        x = torch.cat((cls_tokens, dist_token, x), dim=1)
-    else:
-        cls_tokens = self.cls_token.expand(
-            B, -1, -1
-        )  # stole cls_tokens impl from Phil Wang, thanks
-        x = torch.cat((cls_tokens, x), dim=1)
-
-    x = x + pos_embed
-    x = self.pos_drop(x)
-
-    for blk in self.blocks:
-        x = blk(x)
-
-    x = self.norm(x)
-
-    return x
-
-
-activations = {}
-
-
-def get_activation(name):
-    def hook(model, input, output):
-        activations[name] = output
-
-    return hook
-
-
-def get_readout_oper(vit_features, features, use_readout, start_index=1):
-    if use_readout == "ignore":
-        readout_oper = [Slice(start_index)] * len(features)
-    elif use_readout == "add":
-        readout_oper = [AddReadout(start_index)] * len(features)
-    elif use_readout == "project":
-        readout_oper = [
-            ProjectReadout(vit_features, start_index) for out_feat in features
-        ]
-    else:
-        assert (
-            False
-        ), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'"
-
-    return readout_oper
-
-
-def _make_vit_b16_backbone(
-    model,
-    features=[96, 192, 384, 768],
-    size=[384, 384],
-    hooks=[2, 5, 8, 11],
-    vit_features=768,
-    use_readout="ignore",
-    start_index=1,
-):
-    pretrained = nn.Module()
-
-    pretrained.model = model
-    pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
-    pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
-    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))
-    pretrained.model.blocks[hooks[3]].register_forward_hook(get_activation("4"))
-
-    pretrained.activations = activations
-
-    readout_oper = get_readout_oper(vit_features, features, use_readout, start_index)
-
-    # 32, 48, 136, 384
-    pretrained.act_postprocess1 = nn.Sequential(
-        readout_oper[0],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[0],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-        nn.ConvTranspose2d(
-            in_channels=features[0],
-            out_channels=features[0],
-            kernel_size=4,
-            stride=4,
-            padding=0,
-            bias=True,
-            dilation=1,
-            groups=1,
-        ),
-    )
-
-    pretrained.act_postprocess2 = nn.Sequential(
-        readout_oper[1],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[1],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-        nn.ConvTranspose2d(
-            in_channels=features[1],
-            out_channels=features[1],
-            kernel_size=2,
-            stride=2,
-            padding=0,
-            bias=True,
-            dilation=1,
-            groups=1,
-        ),
-    )
-
-    pretrained.act_postprocess3 = nn.Sequential(
-        readout_oper[2],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[2],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-    )
-
-    pretrained.act_postprocess4 = nn.Sequential(
-        readout_oper[3],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[3],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-        nn.Conv2d(
-            in_channels=features[3],
-            out_channels=features[3],
-            kernel_size=3,
-            stride=2,
-            padding=1,
-        ),
-    )
-
-    pretrained.model.start_index = start_index
-    pretrained.model.patch_size = [16, 16]
-
-    # We inject this function into the VisionTransformer instances so that
-    # we can use it with interpolated position embeddings without modifying the library source.
-    pretrained.model.forward_flex = types.MethodType(forward_flex, pretrained.model)
-    pretrained.model._resize_pos_embed = types.MethodType(
-        _resize_pos_embed, pretrained.model
-    )
-
-    return pretrained
-
-
-def _make_pretrained_vitl16_384(pretrained, use_readout="ignore", hooks=None):
-    model = timm.create_model("vit_large_patch16_384", pretrained=pretrained)
-
-    hooks = [5, 11, 17, 23] if hooks == None else hooks
-    return _make_vit_b16_backbone(
-        model,
-        features=[256, 512, 1024, 1024],
-        hooks=hooks,
-        vit_features=1024,
-        use_readout=use_readout,
-    )
-
-
-def _make_pretrained_vitb16_384(pretrained, use_readout="ignore", hooks=None):
-    model = timm.create_model("vit_base_patch16_384", pretrained=pretrained)
-
-    hooks = [2, 5, 8, 11] if hooks == None else hooks
-    return _make_vit_b16_backbone(
-        model, features=[96, 192, 384, 768], hooks=hooks, use_readout=use_readout
-    )
-
-
-def _make_pretrained_deitb16_384(pretrained, use_readout="ignore", hooks=None):
-    model = timm.create_model("vit_deit_base_patch16_384", pretrained=pretrained)
-
-    hooks = [2, 5, 8, 11] if hooks == None else hooks
-    return _make_vit_b16_backbone(
-        model, features=[96, 192, 384, 768], hooks=hooks, use_readout=use_readout
-    )
-
-
-def _make_pretrained_deitb16_distil_384(pretrained, use_readout="ignore", hooks=None):
-    model = timm.create_model(
-        "vit_deit_base_distilled_patch16_384", pretrained=pretrained
-    )
-
-    hooks = [2, 5, 8, 11] if hooks == None else hooks
-    return _make_vit_b16_backbone(
-        model,
-        features=[96, 192, 384, 768],
-        hooks=hooks,
-        use_readout=use_readout,
-        start_index=2,
-    )
-
-
-def _make_vit_b_rn50_backbone(
-    model,
-    features=[256, 512, 768, 768],
-    size=[384, 384],
-    hooks=[0, 1, 8, 11],
-    vit_features=768,
-    use_vit_only=False,
-    use_readout="ignore",
-    start_index=1,
-):
-    pretrained = nn.Module()
-
-    pretrained.model = model
-
-    if use_vit_only == True:
-        pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
-        pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
-    else:
-        pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
-            get_activation("1")
-        )
-        pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
-            get_activation("2")
-        )
-
-    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))
-    pretrained.model.blocks[hooks[3]].register_forward_hook(get_activation("4"))
-
-    pretrained.activations = activations
-
-    readout_oper = get_readout_oper(vit_features, features, use_readout, start_index)
-
-    if use_vit_only == True:
-        pretrained.act_postprocess1 = nn.Sequential(
-            readout_oper[0],
-            Transpose(1, 2),
-            nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-            nn.Conv2d(
-                in_channels=vit_features,
-                out_channels=features[0],
-                kernel_size=1,
-                stride=1,
-                padding=0,
-            ),
-            nn.ConvTranspose2d(
-                in_channels=features[0],
-                out_channels=features[0],
-                kernel_size=4,
-                stride=4,
-                padding=0,
-                bias=True,
-                dilation=1,
-                groups=1,
-            ),
-        )
-
-        pretrained.act_postprocess2 = nn.Sequential(
-            readout_oper[1],
-            Transpose(1, 2),
-            nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-            nn.Conv2d(
-                in_channels=vit_features,
-                out_channels=features[1],
-                kernel_size=1,
-                stride=1,
-                padding=0,
-            ),
-            nn.ConvTranspose2d(
-                in_channels=features[1],
-                out_channels=features[1],
-                kernel_size=2,
-                stride=2,
-                padding=0,
-                bias=True,
-                dilation=1,
-                groups=1,
-            ),
-        )
-    else:
-        pretrained.act_postprocess1 = nn.Sequential(
-            nn.Identity(), nn.Identity(), nn.Identity()
-        )
-        pretrained.act_postprocess2 = nn.Sequential(
-            nn.Identity(), nn.Identity(), nn.Identity()
-        )
-
-    pretrained.act_postprocess3 = nn.Sequential(
-        readout_oper[2],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[2],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-    )
-
-    pretrained.act_postprocess4 = nn.Sequential(
-        readout_oper[3],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[3],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-        nn.Conv2d(
-            in_channels=features[3],
-            out_channels=features[3],
-            kernel_size=3,
-            stride=2,
-            padding=1,
-        ),
-    )
-
-    pretrained.model.start_index = start_index
-    pretrained.model.patch_size = [16, 16]
-
-    # We inject this function into the VisionTransformer instances so that
-    # we can use it with interpolated position embeddings without modifying the library source.
-    pretrained.model.forward_flex = types.MethodType(forward_flex, pretrained.model)
-
-    # We inject this function into the VisionTransformer instances so that
-    # we can use it with interpolated position embeddings without modifying the library source.
-    pretrained.model._resize_pos_embed = types.MethodType(
-        _resize_pos_embed, pretrained.model
-    )
-
-    return pretrained
-
-
-def _make_pretrained_vitb_rn50_384(
-    pretrained, use_readout="ignore", hooks=None, use_vit_only=False
-):
-    model = timm.create_model("vit_base_resnet50_384", pretrained=pretrained)
-
-    hooks = [0, 1, 8, 11] if hooks == None else hooks
-    return _make_vit_b_rn50_backbone(
-        model,
-        features=[256, 512, 768, 768],
-        size=[384, 384],
-        hooks=hooks,
-        use_vit_only=use_vit_only,
-        use_readout=use_readout,
-    )

controlnet_aux_local/midas/utils.py

@@ -1,189 +0,0 @@
-"""Utils for monoDepth."""
-import sys
-import re
-import numpy as np
-import cv2
-import torch
-
-
-def read_pfm(path):
-    """Read pfm file.
-
-    Args:
-        path (str): path to file
-
-    Returns:
-        tuple: (data, scale)
-    """
-    with open(path, "rb") as file:
-
-        color = None
-        width = None
-        height = None
-        scale = None
-        endian = None
-
-        header = file.readline().rstrip()
-        if header.decode("ascii") == "PF":
-            color = True
-        elif header.decode("ascii") == "Pf":
-            color = False
-        else:
-            raise Exception("Not a PFM file: " + path)
-
-        dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
-        if dim_match:
-            width, height = list(map(int, dim_match.groups()))
-        else:
-            raise Exception("Malformed PFM header.")
-
-        scale = float(file.readline().decode("ascii").rstrip())
-        if scale < 0:
-            # little-endian
-            endian = "<"
-            scale = -scale
-        else:
-            # big-endian
-            endian = ">"
-
-        data = np.fromfile(file, endian + "f")
-        shape = (height, width, 3) if color else (height, width)
-
-        data = np.reshape(data, shape)
-        data = np.flipud(data)
-
-        return data, scale
-
-
-def write_pfm(path, image, scale=1):
-    """Write pfm file.
-
-    Args:
-        path (str): pathto file
-        image (array): data
-        scale (int, optional): Scale. Defaults to 1.
-    """
-
-    with open(path, "wb") as file:
-        color = None
-
-        if image.dtype.name != "float32":
-            raise Exception("Image dtype must be float32.")
-
-        image = np.flipud(image)
-
-        if len(image.shape) == 3 and image.shape[2] == 3:  # color image
-            color = True
-        elif (
-            len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
-        ):  # greyscale
-            color = False
-        else:
-            raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
-
-        file.write("PF\n" if color else "Pf\n".encode())
-        file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
-
-        endian = image.dtype.byteorder
-
-        if endian == "<" or endian == "=" and sys.byteorder == "little":
-            scale = -scale
-
-        file.write("%f\n".encode() % scale)
-
-        image.tofile(file)
-
-
-def read_image(path):
-    """Read image and output RGB image (0-1).
-
-    Args:
-        path (str): path to file
-
-    Returns:
-        array: RGB image (0-1)
-    """
-    img = cv2.imread(path)
-
-    if img.ndim == 2:
-        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
-
-    return img
-
-
-def resize_image(img):
-    """Resize image and make it fit for network.
-
-    Args:
-        img (array): image
-
-    Returns:
-        tensor: data ready for network
-    """
-    height_orig = img.shape[0]
-    width_orig = img.shape[1]
-
-    if width_orig > height_orig:
-        scale = width_orig / 384
-    else:
-        scale = height_orig / 384
-
-    height = (np.ceil(height_orig / scale / 32) * 32).astype(int)
-    width = (np.ceil(width_orig / scale / 32) * 32).astype(int)
-
-    img_resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
-
-    img_resized = (
-        torch.from_numpy(np.transpose(img_resized, (2, 0, 1))).contiguous().float()
-    )
-    img_resized = img_resized.unsqueeze(0)
-
-    return img_resized
-
-
-def resize_depth(depth, width, height):
-    """Resize depth map and bring to CPU (numpy).
-
-    Args:
-        depth (tensor): depth
-        width (int): image width
-        height (int): image height
-
-    Returns:
-        array: processed depth
-    """
-    depth = torch.squeeze(depth[0, :, :, :]).to("cpu")
-
-    depth_resized = cv2.resize(
-        depth.numpy(), (width, height), interpolation=cv2.INTER_CUBIC
-    )
-
-    return depth_resized
-
-def write_depth(path, depth, bits=1):
-    """Write depth map to pfm and png file.
-
-    Args:
-        path (str): filepath without extension
-        depth (array): depth
-    """
-    write_pfm(path + ".pfm", depth.astype(np.float32))
-
-    depth_min = depth.min()
-    depth_max = depth.max()
-
-    max_val = (2**(8*bits))-1
-
-    if depth_max - depth_min > np.finfo("float").eps:
-        out = max_val * (depth - depth_min) / (depth_max - depth_min)
-    else:
-        out = np.zeros(depth.shape, dtype=depth.type)
-
-    if bits == 1:
-        cv2.imwrite(path + ".png", out.astype("uint8"))
-    elif bits == 2:
-        cv2.imwrite(path + ".png", out.astype("uint16"))
-
-    return

controlnet_aux_local/mlsd/__init__.py

@@ -1,79 +0,0 @@
-import os
-import warnings
-
-import cv2
-import numpy as np
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from ..util import HWC3, resize_image
-from .models.mbv2_mlsd_large import MobileV2_MLSD_Large
-from .utils import pred_lines
-
-
-class MLSDdetector:
-    def __init__(self, model):
-        self.model = model
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None, local_files_only=False):
-        if pretrained_model_or_path == "lllyasviel/ControlNet":
-            filename = filename or "annotator/ckpts/mlsd_large_512_fp32.pth"
-        else:
-            filename = filename or "mlsd_large_512_fp32.pth"
-
-        if os.path.isdir(pretrained_model_or_path):
-            model_path = os.path.join(pretrained_model_or_path, filename)
-        else:
-            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir, local_files_only=local_files_only)
-
-        model = MobileV2_MLSD_Large()
-        model.load_state_dict(torch.load(model_path), strict=True)
-        model.eval()
-
-        return cls(model)
-
-    def to(self, device):
-        self.model.to(device)
-        return self
-    
-    def __call__(self, input_image, thr_v=0.1, thr_d=0.1, detect_resolution=512, image_resolution=512, output_type="pil", **kwargs):
-        if "return_pil" in kwargs:
-            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
-            output_type = "pil" if kwargs["return_pil"] else "np"
-        if type(output_type) is bool:
-            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
-            if output_type:
-                output_type = "pil"
-
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-
-        assert input_image.ndim == 3
-        img = input_image
-        img_output = np.zeros_like(img)
-        try:
-            with torch.no_grad():
-                lines = pred_lines(img, self.model, [img.shape[0], img.shape[1]], thr_v, thr_d)
-                for line in lines:
-                    x_start, y_start, x_end, y_end = [int(val) for val in line]
-                    cv2.line(img_output, (x_start, y_start), (x_end, y_end), [255, 255, 255], 1)
-        except Exception as e:
-            pass
-
-        detected_map = img_output[:, :, 0]
-        detected_map = HWC3(detected_map)
-
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-
-        if output_type == "pil":
-            detected_map = Image.fromarray(detected_map)
-
-        return detected_map

controlnet_aux_local/mlsd/models/mbv2_mlsd_large.py

@@ -1,292 +0,0 @@
-import os
-import sys
-import torch
-import torch.nn as nn
-import torch.utils.model_zoo as model_zoo
-from  torch.nn import  functional as F
-
-
-class BlockTypeA(nn.Module):
-    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):
-        super(BlockTypeA, self).__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(in_c2, out_c2, kernel_size=1),
-            nn.BatchNorm2d(out_c2),
-            nn.ReLU(inplace=True)
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_c1, out_c1, kernel_size=1),
-            nn.BatchNorm2d(out_c1),
-            nn.ReLU(inplace=True)
-        )
-        self.upscale = upscale
-
-    def forward(self, a, b):
-        b = self.conv1(b)
-        a = self.conv2(a)
-        if self.upscale:
-             b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)
-        return torch.cat((a, b), dim=1)
-
-
-class BlockTypeB(nn.Module):
-    def __init__(self, in_c, out_c):
-        super(BlockTypeB, self).__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
-            nn.BatchNorm2d(in_c),
-            nn.ReLU()
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
-            nn.BatchNorm2d(out_c),
-            nn.ReLU()
-        )
-
-    def forward(self, x):
-        x = self.conv1(x) + x
-        x = self.conv2(x)
-        return x
-
-class BlockTypeC(nn.Module):
-    def __init__(self, in_c, out_c):
-        super(BlockTypeC, self).__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),
-            nn.BatchNorm2d(in_c),
-            nn.ReLU()
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
-            nn.BatchNorm2d(in_c),
-            nn.ReLU()
-        )
-        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = self.conv3(x)
-        return x
-
-def _make_divisible(v, divisor, min_value=None):
-    """
-    This function is taken from the original tf repo.
-    It ensures that all layers have a channel number that is divisible by 8
-    It can be seen here:
-    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
-    :param v:
-    :param divisor:
-    :param min_value:
-    :return:
-    """
-    if min_value is None:
-        min_value = divisor
-    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
-    # Make sure that round down does not go down by more than 10%.
-    if new_v < 0.9 * v:
-        new_v += divisor
-    return new_v
-
-
-class ConvBNReLU(nn.Sequential):
-    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
-        self.channel_pad = out_planes - in_planes
-        self.stride = stride
-        #padding = (kernel_size - 1) // 2
-
-        # TFLite uses slightly different padding than PyTorch
-        if stride == 2:
-            padding = 0
-        else:
-            padding = (kernel_size - 1) // 2
-
-        super(ConvBNReLU, self).__init__(
-            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
-            nn.BatchNorm2d(out_planes),
-            nn.ReLU6(inplace=True)
-        )
-        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
-
-
-    def forward(self, x):
-        # TFLite uses  different padding
-        if self.stride == 2:
-            x = F.pad(x, (0, 1, 0, 1), "constant", 0)
-            #print(x.shape)
-
-        for module in self:
-            if not isinstance(module, nn.MaxPool2d):
-                x = module(x)
-        return x
-
-
-class InvertedResidual(nn.Module):
-    def __init__(self, inp, oup, stride, expand_ratio):
-        super(InvertedResidual, self).__init__()
-        self.stride = stride
-        assert stride in [1, 2]
-
-        hidden_dim = int(round(inp * expand_ratio))
-        self.use_res_connect = self.stride == 1 and inp == oup
-
-        layers = []
-        if expand_ratio != 1:
-            # pw
-            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
-        layers.extend([
-            # dw
-            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
-            # pw-linear
-            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
-            nn.BatchNorm2d(oup),
-        ])
-        self.conv = nn.Sequential(*layers)
-
-    def forward(self, x):
-        if self.use_res_connect:
-            return x + self.conv(x)
-        else:
-            return self.conv(x)
-
-
-class MobileNetV2(nn.Module):
-    def __init__(self, pretrained=True):
-        """
-        MobileNet V2 main class
-        Args:
-            num_classes (int): Number of classes
-            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
-            inverted_residual_setting: Network structure
-            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
-            Set to 1 to turn off rounding
-            block: Module specifying inverted residual building block for mobilenet
-        """
-        super(MobileNetV2, self).__init__()
-
-        block = InvertedResidual
-        input_channel = 32
-        last_channel = 1280
-        width_mult = 1.0
-        round_nearest = 8
-
-        inverted_residual_setting = [
-            # t, c, n, s
-            [1, 16, 1, 1],
-            [6, 24, 2, 2],
-            [6, 32, 3, 2],
-            [6, 64, 4, 2],
-            [6, 96, 3, 1],
-            #[6, 160, 3, 2],
-            #[6, 320, 1, 1],
-        ]
-
-        # only check the first element, assuming user knows t,c,n,s are required
-        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
-            raise ValueError("inverted_residual_setting should be non-empty "
-                             "or a 4-element list, got {}".format(inverted_residual_setting))
-
-        # building first layer
-        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
-        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
-        features = [ConvBNReLU(4, input_channel, stride=2)]
-        # building inverted residual blocks
-        for t, c, n, s in inverted_residual_setting:
-            output_channel = _make_divisible(c * width_mult, round_nearest)
-            for i in range(n):
-                stride = s if i == 0 else 1
-                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
-                input_channel = output_channel
-
-        self.features = nn.Sequential(*features)
-        self.fpn_selected = [1, 3, 6, 10, 13]
-        # weight initialization
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                nn.init.kaiming_normal_(m.weight, mode='fan_out')
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-            elif isinstance(m, nn.BatchNorm2d):
-                nn.init.ones_(m.weight)
-                nn.init.zeros_(m.bias)
-            elif isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, 0, 0.01)
-                nn.init.zeros_(m.bias)
-        if pretrained:
-           self._load_pretrained_model()
-
-    def _forward_impl(self, x):
-        # This exists since TorchScript doesn't support inheritance, so the superclass method
-        # (this one) needs to have a name other than `forward` that can be accessed in a subclass
-        fpn_features = []
-        for i, f in enumerate(self.features):
-            if i > self.fpn_selected[-1]:
-                break
-            x = f(x)
-            if i in self.fpn_selected:
-                fpn_features.append(x)
-
-        c1, c2, c3, c4, c5 = fpn_features
-        return c1, c2, c3, c4, c5
-
-
-    def forward(self, x):
-        return self._forward_impl(x)
-
-    def _load_pretrained_model(self):
-        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
-        model_dict = {}
-        state_dict = self.state_dict()
-        for k, v in pretrain_dict.items():
-            if k in state_dict:
-                model_dict[k] = v
-        state_dict.update(model_dict)
-        self.load_state_dict(state_dict)
-
-
-class MobileV2_MLSD_Large(nn.Module):
-    def __init__(self):
-        super(MobileV2_MLSD_Large, self).__init__()
-
-        self.backbone = MobileNetV2(pretrained=False)
-        ## A, B
-        self.block15 = BlockTypeA(in_c1= 64, in_c2= 96,
-                                  out_c1= 64, out_c2=64,
-                                  upscale=False)
-        self.block16 = BlockTypeB(128, 64)
-
-        ## A, B
-        self.block17 = BlockTypeA(in_c1 = 32,  in_c2 = 64,
-                                  out_c1= 64,  out_c2= 64)
-        self.block18 = BlockTypeB(128, 64)
-
-        ## A, B
-        self.block19 = BlockTypeA(in_c1=24, in_c2=64,
-                                  out_c1=64, out_c2=64)
-        self.block20 = BlockTypeB(128, 64)
-
-        ## A, B, C
-        self.block21 = BlockTypeA(in_c1=16, in_c2=64,
-                                  out_c1=64, out_c2=64)
-        self.block22 = BlockTypeB(128, 64)
-
-        self.block23 = BlockTypeC(64, 16)
-
-    def forward(self, x):
-        c1, c2, c3, c4, c5 = self.backbone(x)
-
-        x = self.block15(c4, c5)
-        x = self.block16(x)
-
-        x = self.block17(c3, x)
-        x = self.block18(x)
-
-        x = self.block19(c2, x)
-        x = self.block20(x)
-
-        x = self.block21(c1, x)
-        x = self.block22(x)
-        x = self.block23(x)
-        x = x[:, 7:, :, :]
-
-        return x

controlnet_aux_local/mlsd/models/mbv2_mlsd_tiny.py

@@ -1,275 +0,0 @@
-import os
-import sys
-import torch
-import torch.nn as nn
-import torch.utils.model_zoo as model_zoo
-from  torch.nn import  functional as F
-
-
-class BlockTypeA(nn.Module):
-    def __init__(self, in_c1, in_c2, out_c1, out_c2, upscale = True):
-        super(BlockTypeA, self).__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(in_c2, out_c2, kernel_size=1),
-            nn.BatchNorm2d(out_c2),
-            nn.ReLU(inplace=True)
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_c1, out_c1, kernel_size=1),
-            nn.BatchNorm2d(out_c1),
-            nn.ReLU(inplace=True)
-        )
-        self.upscale = upscale
-
-    def forward(self, a, b):
-        b = self.conv1(b)
-        a = self.conv2(a)
-        b = F.interpolate(b, scale_factor=2.0, mode='bilinear', align_corners=True)
-        return torch.cat((a, b), dim=1)
-
-
-class BlockTypeB(nn.Module):
-    def __init__(self, in_c, out_c):
-        super(BlockTypeB, self).__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
-            nn.BatchNorm2d(in_c),
-            nn.ReLU()
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
-            nn.BatchNorm2d(out_c),
-            nn.ReLU()
-        )
-
-    def forward(self, x):
-        x = self.conv1(x) + x
-        x = self.conv2(x)
-        return x
-
-class BlockTypeC(nn.Module):
-    def __init__(self, in_c, out_c):
-        super(BlockTypeC, self).__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=5, dilation=5),
-            nn.BatchNorm2d(in_c),
-            nn.ReLU()
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_c, in_c,  kernel_size=3, padding=1),
-            nn.BatchNorm2d(in_c),
-            nn.ReLU()
-        )
-        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = self.conv3(x)
-        return x
-
-def _make_divisible(v, divisor, min_value=None):
-    """
-    This function is taken from the original tf repo.
-    It ensures that all layers have a channel number that is divisible by 8
-    It can be seen here:
-    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
-    :param v:
-    :param divisor:
-    :param min_value:
-    :return:
-    """
-    if min_value is None:
-        min_value = divisor
-    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
-    # Make sure that round down does not go down by more than 10%.
-    if new_v < 0.9 * v:
-        new_v += divisor
-    return new_v
-
-
-class ConvBNReLU(nn.Sequential):
-    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
-        self.channel_pad = out_planes - in_planes
-        self.stride = stride
-        #padding = (kernel_size - 1) // 2
-
-        # TFLite uses slightly different padding than PyTorch
-        if stride == 2:
-            padding = 0
-        else:
-            padding = (kernel_size - 1) // 2
-
-        super(ConvBNReLU, self).__init__(
-            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
-            nn.BatchNorm2d(out_planes),
-            nn.ReLU6(inplace=True)
-        )
-        self.max_pool = nn.MaxPool2d(kernel_size=stride, stride=stride)
-
-
-    def forward(self, x):
-        # TFLite uses  different padding
-        if self.stride == 2:
-            x = F.pad(x, (0, 1, 0, 1), "constant", 0)
-            #print(x.shape)
-
-        for module in self:
-            if not isinstance(module, nn.MaxPool2d):
-                x = module(x)
-        return x
-
-
-class InvertedResidual(nn.Module):
-    def __init__(self, inp, oup, stride, expand_ratio):
-        super(InvertedResidual, self).__init__()
-        self.stride = stride
-        assert stride in [1, 2]
-
-        hidden_dim = int(round(inp * expand_ratio))
-        self.use_res_connect = self.stride == 1 and inp == oup
-
-        layers = []
-        if expand_ratio != 1:
-            # pw
-            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
-        layers.extend([
-            # dw
-            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
-            # pw-linear
-            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
-            nn.BatchNorm2d(oup),
-        ])
-        self.conv = nn.Sequential(*layers)
-
-    def forward(self, x):
-        if self.use_res_connect:
-            return x + self.conv(x)
-        else:
-            return self.conv(x)
-
-
-class MobileNetV2(nn.Module):
-    def __init__(self, pretrained=True):
-        """
-        MobileNet V2 main class
-        Args:
-            num_classes (int): Number of classes
-            width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
-            inverted_residual_setting: Network structure
-            round_nearest (int): Round the number of channels in each layer to be a multiple of this number
-            Set to 1 to turn off rounding
-            block: Module specifying inverted residual building block for mobilenet
-        """
-        super(MobileNetV2, self).__init__()
-
-        block = InvertedResidual
-        input_channel = 32
-        last_channel = 1280
-        width_mult = 1.0
-        round_nearest = 8
-
-        inverted_residual_setting = [
-            # t, c, n, s
-            [1, 16, 1, 1],
-            [6, 24, 2, 2],
-            [6, 32, 3, 2],
-            [6, 64, 4, 2],
-            #[6, 96, 3, 1],
-            #[6, 160, 3, 2],
-            #[6, 320, 1, 1],
-        ]
-
-        # only check the first element, assuming user knows t,c,n,s are required
-        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
-            raise ValueError("inverted_residual_setting should be non-empty "
-                             "or a 4-element list, got {}".format(inverted_residual_setting))
-
-        # building first layer
-        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
-        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
-        features = [ConvBNReLU(4, input_channel, stride=2)]
-        # building inverted residual blocks
-        for t, c, n, s in inverted_residual_setting:
-            output_channel = _make_divisible(c * width_mult, round_nearest)
-            for i in range(n):
-                stride = s if i == 0 else 1
-                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
-                input_channel = output_channel
-        self.features = nn.Sequential(*features)
-
-        self.fpn_selected = [3, 6, 10]
-        # weight initialization
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                nn.init.kaiming_normal_(m.weight, mode='fan_out')
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-            elif isinstance(m, nn.BatchNorm2d):
-                nn.init.ones_(m.weight)
-                nn.init.zeros_(m.bias)
-            elif isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, 0, 0.01)
-                nn.init.zeros_(m.bias)
-
-        #if pretrained:
-        #    self._load_pretrained_model()
-
-    def _forward_impl(self, x):
-        # This exists since TorchScript doesn't support inheritance, so the superclass method
-        # (this one) needs to have a name other than `forward` that can be accessed in a subclass
-        fpn_features = []
-        for i, f in enumerate(self.features):
-            if i > self.fpn_selected[-1]:
-                break
-            x = f(x)
-            if i in self.fpn_selected:
-                fpn_features.append(x)
-
-        c2, c3, c4 = fpn_features
-        return c2, c3, c4
-
-
-    def forward(self, x):
-        return self._forward_impl(x)
-
-    def _load_pretrained_model(self):
-        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/mobilenet_v2-b0353104.pth')
-        model_dict = {}
-        state_dict = self.state_dict()
-        for k, v in pretrain_dict.items():
-            if k in state_dict:
-                model_dict[k] = v
-        state_dict.update(model_dict)
-        self.load_state_dict(state_dict)
-
-
-class MobileV2_MLSD_Tiny(nn.Module):
-    def __init__(self):
-        super(MobileV2_MLSD_Tiny, self).__init__()
-
-        self.backbone = MobileNetV2(pretrained=True)
-
-        self.block12 = BlockTypeA(in_c1= 32, in_c2= 64,
-                                  out_c1= 64, out_c2=64)
-        self.block13 = BlockTypeB(128, 64)
-
-        self.block14 = BlockTypeA(in_c1 = 24,  in_c2 = 64,
-                                  out_c1= 32,  out_c2= 32)
-        self.block15 = BlockTypeB(64, 64)
-
-        self.block16 = BlockTypeC(64, 16)
-
-    def forward(self, x):
-        c2, c3, c4 = self.backbone(x)
-
-        x = self.block12(c3, c4)
-        x = self.block13(x)
-        x = self.block14(c2, x)
-        x = self.block15(x)
-        x = self.block16(x)
-        x = x[:, 7:, :, :]
-        #print(x.shape)
-        x = F.interpolate(x, scale_factor=2.0, mode='bilinear', align_corners=True)
-
-        return x

controlnet_aux_local/mlsd/utils.py

@@ -1,584 +0,0 @@
-'''
-modified by  lihaoweicv
-pytorch version
-'''
-
-'''
-M-LSD
-Copyright 2021-present NAVER Corp.
-Apache License v2.0
-'''
-
-import os
-import numpy as np
-import cv2
-import torch
-from  torch.nn import functional as F
-
-
-def deccode_output_score_and_ptss(tpMap, topk_n = 200, ksize = 5):
-    '''
-    tpMap:
-    center: tpMap[1, 0, :, :]
-    displacement: tpMap[1, 1:5, :, :]
-    '''
-    b, c, h, w = tpMap.shape
-    assert  b==1, 'only support bsize==1'
-    displacement = tpMap[:, 1:5, :, :][0]
-    center = tpMap[:, 0, :, :]
-    heat = torch.sigmoid(center)
-    hmax = F.max_pool2d( heat, (ksize, ksize), stride=1, padding=(ksize-1)//2)
-    keep = (hmax == heat).float()
-    heat = heat * keep
-    heat = heat.reshape(-1, )
-
-    scores, indices = torch.topk(heat, topk_n, dim=-1, largest=True)
-    yy = torch.floor_divide(indices, w).unsqueeze(-1)
-    xx = torch.fmod(indices, w).unsqueeze(-1)
-    ptss = torch.cat((yy, xx),dim=-1)
-
-    ptss   = ptss.detach().cpu().numpy()
-    scores = scores.detach().cpu().numpy()
-    displacement = displacement.detach().cpu().numpy()
-    displacement = displacement.transpose((1,2,0))
-    return  ptss, scores, displacement
-
-
-def pred_lines(image, model,
-               input_shape=[512, 512],
-               score_thr=0.10,
-               dist_thr=20.0):
-    h, w, _ = image.shape
-
-    device = next(iter(model.parameters())).device
-    h_ratio, w_ratio = [h / input_shape[0], w / input_shape[1]]
-
-    resized_image = np.concatenate([cv2.resize(image, (input_shape[1], input_shape[0]), interpolation=cv2.INTER_AREA),
-                                    np.ones([input_shape[0], input_shape[1], 1])], axis=-1)
-
-    resized_image = resized_image.transpose((2,0,1))
-    batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
-    batch_image = (batch_image / 127.5) - 1.0
-
-    batch_image = torch.from_numpy(batch_image).float()
-    batch_image = batch_image.to(device)
-    outputs = model(batch_image)
-    pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
-    start = vmap[:, :, :2]
-    end = vmap[:, :, 2:]
-    dist_map = np.sqrt(np.sum((start - end) ** 2, axis=-1))
-
-    segments_list = []
-    for center, score in zip(pts, pts_score):
-        y, x = center
-        distance = dist_map[y, x]
-        if score > score_thr and distance > dist_thr:
-            disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
-            x_start = x + disp_x_start
-            y_start = y + disp_y_start
-            x_end = x + disp_x_end
-            y_end = y + disp_y_end
-            segments_list.append([x_start, y_start, x_end, y_end])
-
-    lines = 2 * np.array(segments_list)  # 256 > 512
-    lines[:, 0] = lines[:, 0] * w_ratio
-    lines[:, 1] = lines[:, 1] * h_ratio
-    lines[:, 2] = lines[:, 2] * w_ratio
-    lines[:, 3] = lines[:, 3] * h_ratio
-
-    return lines
-
-
-def pred_squares(image,
-                 model,
-                 input_shape=[512, 512],
-                 params={'score': 0.06,
-                         'outside_ratio': 0.28,
-                         'inside_ratio': 0.45,
-                         'w_overlap': 0.0,
-                         'w_degree': 1.95,
-                         'w_length': 0.0,
-                         'w_area': 1.86,
-                         'w_center': 0.14}):
-    '''
-    shape = [height, width]
-    '''
-    h, w, _ = image.shape
-    original_shape = [h, w]
-    device = next(iter(model.parameters())).device
-
-    resized_image = np.concatenate([cv2.resize(image, (input_shape[0], input_shape[1]), interpolation=cv2.INTER_AREA),
-                                    np.ones([input_shape[0], input_shape[1], 1])], axis=-1)
-    resized_image = resized_image.transpose((2, 0, 1))
-    batch_image = np.expand_dims(resized_image, axis=0).astype('float32')
-    batch_image = (batch_image / 127.5) - 1.0
-
-    batch_image = torch.from_numpy(batch_image).float().to(device)
-    outputs = model(batch_image)
-
-    pts, pts_score, vmap = deccode_output_score_and_ptss(outputs, 200, 3)
-    start = vmap[:, :, :2]  # (x, y)
-    end = vmap[:, :, 2:]  # (x, y)
-    dist_map = np.sqrt(np.sum((start - end) ** 2, axis=-1))
-
-    junc_list = []
-    segments_list = []
-    for junc, score in zip(pts, pts_score):
-        y, x = junc
-        distance = dist_map[y, x]
-        if score > params['score'] and distance > 20.0:
-            junc_list.append([x, y])
-            disp_x_start, disp_y_start, disp_x_end, disp_y_end = vmap[y, x, :]
-            d_arrow = 1.0
-            x_start = x + d_arrow * disp_x_start
-            y_start = y + d_arrow * disp_y_start
-            x_end = x + d_arrow * disp_x_end
-            y_end = y + d_arrow * disp_y_end
-            segments_list.append([x_start, y_start, x_end, y_end])
-
-    segments = np.array(segments_list)
-
-    ####### post processing for squares
-    # 1. get unique lines
-    point = np.array([[0, 0]])
-    point = point[0]
-    start = segments[:, :2]
-    end = segments[:, 2:]
-    diff = start - end
-    a = diff[:, 1]
-    b = -diff[:, 0]
-    c = a * start[:, 0] + b * start[:, 1]
-
-    d = np.abs(a * point[0] + b * point[1] - c) / np.sqrt(a ** 2 + b ** 2 + 1e-10)
-    theta = np.arctan2(diff[:, 0], diff[:, 1]) * 180 / np.pi
-    theta[theta < 0.0] += 180
-    hough = np.concatenate([d[:, None], theta[:, None]], axis=-1)
-
-    d_quant = 1
-    theta_quant = 2
-    hough[:, 0] //= d_quant
-    hough[:, 1] //= theta_quant
-    _, indices, counts = np.unique(hough, axis=0, return_index=True, return_counts=True)
-
-    acc_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1], dtype='float32')
-    idx_map = np.zeros([512 // d_quant + 1, 360 // theta_quant + 1], dtype='int32') - 1
-    yx_indices = hough[indices, :].astype('int32')
-    acc_map[yx_indices[:, 0], yx_indices[:, 1]] = counts
-    idx_map[yx_indices[:, 0], yx_indices[:, 1]] = indices
-
-    acc_map_np = acc_map
-    # acc_map = acc_map[None, :, :, None]
-    #
-    # ### fast suppression using tensorflow op
-    # acc_map = tf.constant(acc_map, dtype=tf.float32)
-    # max_acc_map = tf.keras.layers.MaxPool2D(pool_size=(5, 5), strides=1, padding='same')(acc_map)
-    # acc_map = acc_map * tf.cast(tf.math.equal(acc_map, max_acc_map), tf.float32)
-    # flatten_acc_map = tf.reshape(acc_map, [1, -1])
-    # topk_values, topk_indices = tf.math.top_k(flatten_acc_map, k=len(pts))
-    # _, h, w, _ = acc_map.shape
-    # y = tf.expand_dims(topk_indices // w, axis=-1)
-    # x = tf.expand_dims(topk_indices % w, axis=-1)
-    # yx = tf.concat([y, x], axis=-1)
-
-    ### fast suppression using pytorch op
-    acc_map = torch.from_numpy(acc_map_np).unsqueeze(0).unsqueeze(0)
-    _,_, h, w = acc_map.shape
-    max_acc_map = F.max_pool2d(acc_map,kernel_size=5, stride=1, padding=2)
-    acc_map = acc_map * ( (acc_map == max_acc_map).float() )
-    flatten_acc_map = acc_map.reshape([-1, ])
-
-    scores, indices = torch.topk(flatten_acc_map, len(pts), dim=-1, largest=True)
-    yy = torch.div(indices, w, rounding_mode='floor').unsqueeze(-1)
-    xx = torch.fmod(indices, w).unsqueeze(-1)
-    yx = torch.cat((yy, xx), dim=-1)
-
-    yx = yx.detach().cpu().numpy()
-
-    topk_values = scores.detach().cpu().numpy()
-    indices = idx_map[yx[:, 0], yx[:, 1]]
-    basis = 5 // 2
-
-    merged_segments = []
-    for yx_pt, max_indice, value in zip(yx, indices, topk_values):
-        y, x = yx_pt
-        if max_indice == -1 or value == 0:
-            continue
-        segment_list = []
-        for y_offset in range(-basis, basis + 1):
-            for x_offset in range(-basis, basis + 1):
-                indice = idx_map[y + y_offset, x + x_offset]
-                cnt = int(acc_map_np[y + y_offset, x + x_offset])
-                if indice != -1:
-                    segment_list.append(segments[indice])
-                if cnt > 1:
-                    check_cnt = 1
-                    current_hough = hough[indice]
-                    for new_indice, new_hough in enumerate(hough):
-                        if (current_hough == new_hough).all() and indice != new_indice:
-                            segment_list.append(segments[new_indice])
-                            check_cnt += 1
-                        if check_cnt == cnt:
-                            break
-        group_segments = np.array(segment_list).reshape([-1, 2])
-        sorted_group_segments = np.sort(group_segments, axis=0)
-        x_min, y_min = sorted_group_segments[0, :]
-        x_max, y_max = sorted_group_segments[-1, :]
-
-        deg = theta[max_indice]
-        if deg >= 90:
-            merged_segments.append([x_min, y_max, x_max, y_min])
-        else:
-            merged_segments.append([x_min, y_min, x_max, y_max])
-
-    # 2. get intersections
-    new_segments = np.array(merged_segments)  # (x1, y1, x2, y2)
-    start = new_segments[:, :2]  # (x1, y1)
-    end = new_segments[:, 2:]  # (x2, y2)
-    new_centers = (start + end) / 2.0
-    diff = start - end
-    dist_segments = np.sqrt(np.sum(diff ** 2, axis=-1))
-
-    # ax + by = c
-    a = diff[:, 1]
-    b = -diff[:, 0]
-    c = a * start[:, 0] + b * start[:, 1]
-    pre_det = a[:, None] * b[None, :]
-    det = pre_det - np.transpose(pre_det)
-
-    pre_inter_y = a[:, None] * c[None, :]
-    inter_y = (pre_inter_y - np.transpose(pre_inter_y)) / (det + 1e-10)
-    pre_inter_x = c[:, None] * b[None, :]
-    inter_x = (pre_inter_x - np.transpose(pre_inter_x)) / (det + 1e-10)
-    inter_pts = np.concatenate([inter_x[:, :, None], inter_y[:, :, None]], axis=-1).astype('int32')
-
-    # 3. get corner information
-    # 3.1 get distance
-    '''
-    dist_segments:
-        | dist(0), dist(1), dist(2), ...|
-    dist_inter_to_segment1:
-        | dist(inter,0), dist(inter,0), dist(inter,0), ... |
-        | dist(inter,1), dist(inter,1), dist(inter,1), ... |
-        ...
-    dist_inter_to_semgnet2:
-        | dist(inter,0), dist(inter,1), dist(inter,2), ... |
-        | dist(inter,0), dist(inter,1), dist(inter,2), ... |
-        ...
-    '''
-
-    dist_inter_to_segment1_start = np.sqrt(
-        np.sum(((inter_pts - start[:, None, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
-    dist_inter_to_segment1_end = np.sqrt(
-        np.sum(((inter_pts - end[:, None, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
-    dist_inter_to_segment2_start = np.sqrt(
-        np.sum(((inter_pts - start[None, :, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
-    dist_inter_to_segment2_end = np.sqrt(
-        np.sum(((inter_pts - end[None, :, :]) ** 2), axis=-1, keepdims=True))  # [n_batch, n_batch, 1]
-
-    # sort ascending
-    dist_inter_to_segment1 = np.sort(
-        np.concatenate([dist_inter_to_segment1_start, dist_inter_to_segment1_end], axis=-1),
-        axis=-1)  # [n_batch, n_batch, 2]
-    dist_inter_to_segment2 = np.sort(
-        np.concatenate([dist_inter_to_segment2_start, dist_inter_to_segment2_end], axis=-1),
-        axis=-1)  # [n_batch, n_batch, 2]
-
-    # 3.2 get degree
-    inter_to_start = new_centers[:, None, :] - inter_pts
-    deg_inter_to_start = np.arctan2(inter_to_start[:, :, 1], inter_to_start[:, :, 0]) * 180 / np.pi
-    deg_inter_to_start[deg_inter_to_start < 0.0] += 360
-    inter_to_end = new_centers[None, :, :] - inter_pts
-    deg_inter_to_end = np.arctan2(inter_to_end[:, :, 1], inter_to_end[:, :, 0]) * 180 / np.pi
-    deg_inter_to_end[deg_inter_to_end < 0.0] += 360
-
-    '''
-    B -- G
-    |    |
-    C -- R
-    B : blue / G: green / C: cyan / R: red
-
-    0 -- 1
-    |    |
-    3 -- 2
-    '''
-    # rename variables
-    deg1_map, deg2_map = deg_inter_to_start, deg_inter_to_end
-    # sort deg ascending
-    deg_sort = np.sort(np.concatenate([deg1_map[:, :, None], deg2_map[:, :, None]], axis=-1), axis=-1)
-
-    deg_diff_map = np.abs(deg1_map - deg2_map)
-    # we only consider the smallest degree of intersect
-    deg_diff_map[deg_diff_map > 180] = 360 - deg_diff_map[deg_diff_map > 180]
-
-    # define available degree range
-    deg_range = [60, 120]
-
-    corner_dict = {corner_info: [] for corner_info in range(4)}
-    inter_points = []
-    for i in range(inter_pts.shape[0]):
-        for j in range(i + 1, inter_pts.shape[1]):
-            # i, j > line index, always i < j
-            x, y = inter_pts[i, j, :]
-            deg1, deg2 = deg_sort[i, j, :]
-            deg_diff = deg_diff_map[i, j]
-
-            check_degree = deg_diff > deg_range[0] and deg_diff < deg_range[1]
-
-            outside_ratio = params['outside_ratio']  # over ratio >>> drop it!
-            inside_ratio = params['inside_ratio']  # over ratio >>> drop it!
-            check_distance = ((dist_inter_to_segment1[i, j, 1] >= dist_segments[i] and \
-                               dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * outside_ratio) or \
-                              (dist_inter_to_segment1[i, j, 1] <= dist_segments[i] and \
-                               dist_inter_to_segment1[i, j, 0] <= dist_segments[i] * inside_ratio)) and \
-                             ((dist_inter_to_segment2[i, j, 1] >= dist_segments[j] and \
-                               dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * outside_ratio) or \
-                              (dist_inter_to_segment2[i, j, 1] <= dist_segments[j] and \
-                               dist_inter_to_segment2[i, j, 0] <= dist_segments[j] * inside_ratio))
-
-            if check_degree and check_distance:
-                corner_info = None
-
-                if (deg1 >= 0 and deg1 <= 45 and deg2 >= 45 and deg2 <= 120) or \
-                        (deg2 >= 315 and deg1 >= 45 and deg1 <= 120):
-                    corner_info, color_info = 0, 'blue'
-                elif (deg1 >= 45 and deg1 <= 125 and deg2 >= 125 and deg2 <= 225):
-                    corner_info, color_info = 1, 'green'
-                elif (deg1 >= 125 and deg1 <= 225 and deg2 >= 225 and deg2 <= 315):
-                    corner_info, color_info = 2, 'black'
-                elif (deg1 >= 0 and deg1 <= 45 and deg2 >= 225 and deg2 <= 315) or \
-                        (deg2 >= 315 and deg1 >= 225 and deg1 <= 315):
-                    corner_info, color_info = 3, 'cyan'
-                else:
-                    corner_info, color_info = 4, 'red'  # we don't use it
-                    continue
-
-                corner_dict[corner_info].append([x, y, i, j])
-                inter_points.append([x, y])
-
-    square_list = []
-    connect_list = []
-    segments_list = []
-    for corner0 in corner_dict[0]:
-        for corner1 in corner_dict[1]:
-            connect01 = False
-            for corner0_line in corner0[2:]:
-                if corner0_line in corner1[2:]:
-                    connect01 = True
-                    break
-            if connect01:
-                for corner2 in corner_dict[2]:
-                    connect12 = False
-                    for corner1_line in corner1[2:]:
-                        if corner1_line in corner2[2:]:
-                            connect12 = True
-                            break
-                    if connect12:
-                        for corner3 in corner_dict[3]:
-                            connect23 = False
-                            for corner2_line in corner2[2:]:
-                                if corner2_line in corner3[2:]:
-                                    connect23 = True
-                                    break
-                            if connect23:
-                                for corner3_line in corner3[2:]:
-                                    if corner3_line in corner0[2:]:
-                                        # SQUARE!!!
-                                        '''
-                                        0 -- 1
-                                        |    |
-                                        3 -- 2
-                                        square_list:
-                                            order: 0 > 1 > 2 > 3
-                                            | x0, y0, x1, y1, x2, y2, x3, y3 |
-                                            | x0, y0, x1, y1, x2, y2, x3, y3 |
-                                            ...
-                                        connect_list:
-                                            order: 01 > 12 > 23 > 30
-                                            | line_idx01, line_idx12, line_idx23, line_idx30 |
-                                            | line_idx01, line_idx12, line_idx23, line_idx30 |
-                                            ...
-                                        segments_list:
-                                            order: 0 > 1 > 2 > 3
-                                            | line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i, line_idx2_j, line_idx3_i, line_idx3_j |
-                                            | line_idx0_i, line_idx0_j, line_idx1_i, line_idx1_j, line_idx2_i, line_idx2_j, line_idx3_i, line_idx3_j |
-                                            ...
-                                        '''
-                                        square_list.append(corner0[:2] + corner1[:2] + corner2[:2] + corner3[:2])
-                                        connect_list.append([corner0_line, corner1_line, corner2_line, corner3_line])
-                                        segments_list.append(corner0[2:] + corner1[2:] + corner2[2:] + corner3[2:])
-
-    def check_outside_inside(segments_info, connect_idx):
-        # return 'outside or inside', min distance, cover_param, peri_param
-        if connect_idx == segments_info[0]:
-            check_dist_mat = dist_inter_to_segment1
-        else:
-            check_dist_mat = dist_inter_to_segment2
-
-        i, j = segments_info
-        min_dist, max_dist = check_dist_mat[i, j, :]
-        connect_dist = dist_segments[connect_idx]
-        if max_dist > connect_dist:
-            return 'outside', min_dist, 0, 1
-        else:
-            return 'inside', min_dist, -1, -1
-
-    top_square = None
-
-    try:
-        map_size = input_shape[0] / 2
-        squares = np.array(square_list).reshape([-1, 4, 2])
-        score_array = []
-        connect_array = np.array(connect_list)
-        segments_array = np.array(segments_list).reshape([-1, 4, 2])
-
-        # get degree of corners:
-        squares_rollup = np.roll(squares, 1, axis=1)
-        squares_rolldown = np.roll(squares, -1, axis=1)
-        vec1 = squares_rollup - squares
-        normalized_vec1 = vec1 / (np.linalg.norm(vec1, axis=-1, keepdims=True) + 1e-10)
-        vec2 = squares_rolldown - squares
-        normalized_vec2 = vec2 / (np.linalg.norm(vec2, axis=-1, keepdims=True) + 1e-10)
-        inner_products = np.sum(normalized_vec1 * normalized_vec2, axis=-1)  # [n_squares, 4]
-        squares_degree = np.arccos(inner_products) * 180 / np.pi  # [n_squares, 4]
-
-        # get square score
-        overlap_scores = []
-        degree_scores = []
-        length_scores = []
-
-        for connects, segments, square, degree in zip(connect_array, segments_array, squares, squares_degree):
-            '''
-            0 -- 1
-            |    |
-            3 -- 2
-
-            # segments: [4, 2]
-            # connects: [4]
-            '''
-
-            ###################################### OVERLAP SCORES
-            cover = 0
-            perimeter = 0
-            # check 0 > 1 > 2 > 3
-            square_length = []
-
-            for start_idx in range(4):
-                end_idx = (start_idx + 1) % 4
-
-                connect_idx = connects[start_idx]  # segment idx of segment01
-                start_segments = segments[start_idx]
-                end_segments = segments[end_idx]
-
-                start_point = square[start_idx]
-                end_point = square[end_idx]
-
-                # check whether outside or inside
-                start_position, start_min, start_cover_param, start_peri_param = check_outside_inside(start_segments,
-                                                                                                      connect_idx)
-                end_position, end_min, end_cover_param, end_peri_param = check_outside_inside(end_segments, connect_idx)
-
-                cover += dist_segments[connect_idx] + start_cover_param * start_min + end_cover_param * end_min
-                perimeter += dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min
-
-                square_length.append(
-                    dist_segments[connect_idx] + start_peri_param * start_min + end_peri_param * end_min)
-
-            overlap_scores.append(cover / perimeter)
-            ######################################
-            ###################################### DEGREE SCORES
-            '''
-            deg0 vs deg2
-            deg1 vs deg3
-            '''
-            deg0, deg1, deg2, deg3 = degree
-            deg_ratio1 = deg0 / deg2
-            if deg_ratio1 > 1.0:
-                deg_ratio1 = 1 / deg_ratio1
-            deg_ratio2 = deg1 / deg3
-            if deg_ratio2 > 1.0:
-                deg_ratio2 = 1 / deg_ratio2
-            degree_scores.append((deg_ratio1 + deg_ratio2) / 2)
-            ######################################
-            ###################################### LENGTH SCORES
-            '''
-            len0 vs len2
-            len1 vs len3
-            '''
-            len0, len1, len2, len3 = square_length
-            len_ratio1 = len0 / len2 if len2 > len0 else len2 / len0
-            len_ratio2 = len1 / len3 if len3 > len1 else len3 / len1
-            length_scores.append((len_ratio1 + len_ratio2) / 2)
-
-            ######################################
-
-        overlap_scores = np.array(overlap_scores)
-        overlap_scores /= np.max(overlap_scores)
-
-        degree_scores = np.array(degree_scores)
-        # degree_scores /= np.max(degree_scores)
-
-        length_scores = np.array(length_scores)
-
-        ###################################### AREA SCORES
-        area_scores = np.reshape(squares, [-1, 4, 2])
-        area_x = area_scores[:, :, 0]
-        area_y = area_scores[:, :, 1]
-        correction = area_x[:, -1] * area_y[:, 0] - area_y[:, -1] * area_x[:, 0]
-        area_scores = np.sum(area_x[:, :-1] * area_y[:, 1:], axis=-1) - np.sum(area_y[:, :-1] * area_x[:, 1:], axis=-1)
-        area_scores = 0.5 * np.abs(area_scores + correction)
-        area_scores /= (map_size * map_size)  # np.max(area_scores)
-        ######################################
-
-        ###################################### CENTER SCORES
-        centers = np.array([[256 // 2, 256 // 2]], dtype='float32')  # [1, 2]
-        # squares: [n, 4, 2]
-        square_centers = np.mean(squares, axis=1)  # [n, 2]
-        center2center = np.sqrt(np.sum((centers - square_centers) ** 2))
-        center_scores = center2center / (map_size / np.sqrt(2.0))
-
-        '''
-        score_w = [overlap, degree, area, center, length]
-        '''
-        score_w = [0.0, 1.0, 10.0, 0.5, 1.0]
-        score_array = params['w_overlap'] * overlap_scores \
-                      + params['w_degree'] * degree_scores \
-                      + params['w_area'] * area_scores \
-                      - params['w_center'] * center_scores \
-                      + params['w_length'] * length_scores
-
-        best_square = []
-
-        sorted_idx = np.argsort(score_array)[::-1]
-        score_array = score_array[sorted_idx]
-        squares = squares[sorted_idx]
-
-    except Exception as e:
-        pass
-
-    '''return list
-    merged_lines, squares, scores
-    '''
-
-    try:
-        new_segments[:, 0] = new_segments[:, 0] * 2 / input_shape[1] * original_shape[1]
-        new_segments[:, 1] = new_segments[:, 1] * 2 / input_shape[0] * original_shape[0]
-        new_segments[:, 2] = new_segments[:, 2] * 2 / input_shape[1] * original_shape[1]
-        new_segments[:, 3] = new_segments[:, 3] * 2 / input_shape[0] * original_shape[0]
-    except:
-        new_segments = []
-
-    try:
-        squares[:, :, 0] = squares[:, :, 0] * 2 / input_shape[1] * original_shape[1]
-        squares[:, :, 1] = squares[:, :, 1] * 2 / input_shape[0] * original_shape[0]
-    except:
-        squares = []
-        score_array = []
-
-    try:
-        inter_points = np.array(inter_points)
-        inter_points[:, 0] = inter_points[:, 0] * 2 / input_shape[1] * original_shape[1]
-        inter_points[:, 1] = inter_points[:, 1] * 2 / input_shape[0] * original_shape[0]
-    except:
-        inter_points = []
-
-    return new_segments, squares, score_array, inter_points

controlnet_aux_local/open_pose/__init__.py

@@ -1,234 +0,0 @@
-# Openpose
-# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
-# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
-# 3rd Edited by ControlNet
-# 4th Edited by ControlNet (added face and correct hands)
-# 5th Edited by ControlNet (Improved JSON serialization/deserialization, and lots of bug fixs)
-# This preprocessor is licensed by CMU for non-commercial use only.
-
-
-import os
-
-os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
-
-import json
-import warnings
-from typing import Callable, List, NamedTuple, Tuple, Union
-
-import cv2
-import numpy as np
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-from ..util import HWC3, resize_image
-from . import util
-from .body import Body, BodyResult, Keypoint
-from .face import Face
-from .hand import Hand
-
-HandResult = List[Keypoint]
-FaceResult = List[Keypoint]
-
-class PoseResult(NamedTuple):
-    body: BodyResult
-    left_hand: Union[HandResult, None]
-    right_hand: Union[HandResult, None]
-    face: Union[FaceResult, None]
-
-def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, draw_face=True):
-    """
-    Draw the detected poses on an empty canvas.
-
-    Args:
-        poses (List[PoseResult]): A list of PoseResult objects containing the detected poses.
-        H (int): The height of the canvas.
-        W (int): The width of the canvas.
-        draw_body (bool, optional): Whether to draw body keypoints. Defaults to True.
-        draw_hand (bool, optional): Whether to draw hand keypoints. Defaults to True.
-        draw_face (bool, optional): Whether to draw face keypoints. Defaults to True.
-
-    Returns:
-        numpy.ndarray: A 3D numpy array representing the canvas with the drawn poses.
-    """
-    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
-
-    for pose in poses:
-        if draw_body:
-            canvas = util.draw_bodypose(canvas, pose.body.keypoints)
-
-        if draw_hand:
-            canvas = util.draw_handpose(canvas, pose.left_hand)
-            canvas = util.draw_handpose(canvas, pose.right_hand)
-
-        if draw_face:
-            canvas = util.draw_facepose(canvas, pose.face)
-
-    return canvas
-    
-    
-class OpenposeDetector:
-    """
-    A class for detecting human poses in images using the Openpose model.
-
-    Attributes:
-        model_dir (str): Path to the directory where the pose models are stored.
-    """
-    def __init__(self, body_estimation, hand_estimation=None, face_estimation=None):
-        self.body_estimation = body_estimation
-        self.hand_estimation = hand_estimation
-        self.face_estimation = face_estimation
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_or_path, filename=None, hand_filename=None, face_filename=None, cache_dir=None, local_files_only=False):
-
-        if pretrained_model_or_path == "lllyasviel/ControlNet":
-            filename = filename or "annotator/ckpts/body_pose_model.pth"
-            hand_filename = hand_filename or "annotator/ckpts/hand_pose_model.pth"
-            face_filename = face_filename or "facenet.pth"
-
-            face_pretrained_model_or_path = "lllyasviel/Annotators"
-        else:
-            filename = filename or "body_pose_model.pth"
-            hand_filename = hand_filename or "hand_pose_model.pth"
-            face_filename = face_filename or "facenet.pth"
-
-            face_pretrained_model_or_path = pretrained_model_or_path
-
-        if os.path.isdir(pretrained_model_or_path):
-            body_model_path = os.path.join(pretrained_model_or_path, filename)
-            hand_model_path = os.path.join(pretrained_model_or_path, hand_filename)
-            face_model_path = os.path.join(face_pretrained_model_or_path, face_filename)
-        else:
-            body_model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir, local_files_only=local_files_only)
-            hand_model_path = hf_hub_download(pretrained_model_or_path, hand_filename, cache_dir=cache_dir, local_files_only=local_files_only)
-            face_model_path = hf_hub_download(face_pretrained_model_or_path, face_filename, cache_dir=cache_dir, local_files_only=local_files_only)
-
-        body_estimation = Body(body_model_path)
-        hand_estimation = Hand(hand_model_path)
-        face_estimation = Face(face_model_path)
-
-        return cls(body_estimation, hand_estimation, face_estimation)
-
-    def to(self, device):
-        self.body_estimation.to(device)
-        self.hand_estimation.to(device)
-        self.face_estimation.to(device)
-        return self
-
-    def detect_hands(self, body: BodyResult, oriImg) -> Tuple[Union[HandResult, None], Union[HandResult, None]]:
-        left_hand = None
-        right_hand = None
-        H, W, _ = oriImg.shape
-        for x, y, w, is_left in util.handDetect(body, oriImg):
-            peaks = self.hand_estimation(oriImg[y:y+w, x:x+w, :]).astype(np.float32)
-            if peaks.ndim == 2 and peaks.shape[1] == 2:
-                peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)
-                peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)
-                
-                hand_result = [
-                    Keypoint(x=peak[0], y=peak[1])
-                    for peak in peaks
-                ]
-
-                if is_left:
-                    left_hand = hand_result
-                else:
-                    right_hand = hand_result
-
-        return left_hand, right_hand
-
-    def detect_face(self, body: BodyResult, oriImg) -> Union[FaceResult, None]:
-        face = util.faceDetect(body, oriImg)
-        if face is None:
-            return None
-        
-        x, y, w = face
-        H, W, _ = oriImg.shape
-        heatmaps = self.face_estimation(oriImg[y:y+w, x:x+w, :])
-        peaks = self.face_estimation.compute_peaks_from_heatmaps(heatmaps).astype(np.float32)
-        if peaks.ndim == 2 and peaks.shape[1] == 2:
-            peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)
-            peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)
-            return [
-                Keypoint(x=peak[0], y=peak[1])
-                for peak in peaks
-            ]
-        
-        return None
-
-    def detect_poses(self, oriImg, include_hand=False, include_face=False) -> List[PoseResult]:
-        """
-        Detect poses in the given image.
-            Args:
-                oriImg (numpy.ndarray): The input image for pose detection.
-                include_hand (bool, optional): Whether to include hand detection. Defaults to False.
-                include_face (bool, optional): Whether to include face detection. Defaults to False.
-
-        Returns:
-            List[PoseResult]: A list of PoseResult objects containing the detected poses.
-        """
-        oriImg = oriImg[:, :, ::-1].copy()
-        H, W, C = oriImg.shape
-        with torch.no_grad():
-            candidate, subset = self.body_estimation(oriImg)
-            bodies = self.body_estimation.format_body_result(candidate, subset)
-
-            results = []
-            for body in bodies:
-                left_hand, right_hand, face = (None,) * 3
-                if include_hand:
-                    left_hand, right_hand = self.detect_hands(body, oriImg)
-                if include_face:
-                    face = self.detect_face(body, oriImg)
-                
-                results.append(PoseResult(BodyResult(
-                    keypoints=[
-                        Keypoint(
-                            x=keypoint.x / float(W),
-                            y=keypoint.y / float(H)
-                        ) if keypoint is not None else None
-                        for keypoint in body.keypoints
-                    ], 
-                    total_score=body.total_score,
-                    total_parts=body.total_parts
-                ), left_hand, right_hand, face))
-            
-            return results
-        
-    def __call__(self, input_image, detect_resolution=512, image_resolution=512, include_body=True, include_hand=False, include_face=False, hand_and_face=None, output_type="pil", **kwargs):
-        if hand_and_face is not None:
-            warnings.warn("hand_and_face is deprecated. Use include_hand and include_face instead.", DeprecationWarning)
-            include_hand = hand_and_face
-            include_face = hand_and_face
-
-        if "return_pil" in kwargs:
-            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
-            output_type = "pil" if kwargs["return_pil"] else "np"
-        if type(output_type) is bool:
-            warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
-            if output_type:
-                output_type = "pil"
-
-        if not isinstance(input_image, np.ndarray):
-            input_image = np.array(input_image, dtype=np.uint8)
-
-        input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
-        H, W, C = input_image.shape
-        
-        poses = self.detect_poses(input_image, include_hand, include_face)
-        canvas = draw_poses(poses, H, W, draw_body=include_body, draw_hand=include_hand, draw_face=include_face) 
-
-        detected_map = canvas
-        detected_map = HWC3(detected_map)
-        
-        img = resize_image(input_image, image_resolution)
-        H, W, C = img.shape
-
-        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-
-        if output_type == "pil":
-            detected_map = Image.fromarray(detected_map)
-
-        return detected_map

controlnet_aux_local/open_pose/body.py

@@ -1,260 +0,0 @@
-import math
-from typing import List, NamedTuple, Union
-
-import cv2
-import numpy as np
-import torch
-from scipy.ndimage.filters import gaussian_filter
-
-from . import util
-from .model import bodypose_model
-
-
-class Keypoint(NamedTuple):
-    x: float
-    y: float
-    score: float = 1.0
-    id: int = -1
-
-
-class BodyResult(NamedTuple):
-    # Note: Using `Union` instead of `|` operator as the ladder is a Python
-    # 3.10 feature.
-    # Annotator code should be Python 3.8 Compatible, as controlnet repo uses
-    # Python 3.8 environment.
-    # https://github.com/lllyasviel/ControlNet/blob/d3284fcd0972c510635a4f5abe2eeb71dc0de524/environment.yaml#L6
-    keypoints: List[Union[Keypoint, None]]
-    total_score: float
-    total_parts: int
-
-
-class Body(object):
-    def __init__(self, model_path):
-        self.model = bodypose_model()
-        model_dict = util.transfer(self.model, torch.load(model_path))
-        self.model.load_state_dict(model_dict)
-        self.model.eval()
-
-    def to(self, device):
-        self.model.to(device)
-        return self
-
-    def __call__(self, oriImg):
-        device = next(iter(self.model.parameters())).device
-        # scale_search = [0.5, 1.0, 1.5, 2.0]
-        scale_search = [0.5]
-        boxsize = 368
-        stride = 8
-        padValue = 128
-        thre1 = 0.1
-        thre2 = 0.05
-        multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
-        heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
-        paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
-
-        for m in range(len(multiplier)):
-            scale = multiplier[m]
-            imageToTest = util.smart_resize_k(oriImg, fx=scale, fy=scale)
-            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
-            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
-            im = np.ascontiguousarray(im)
-
-            data = torch.from_numpy(im).float()
-            data = data.to(device)
-            # data = data.permute([2, 0, 1]).unsqueeze(0).float()
-            with torch.no_grad():
-                Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
-            Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
-            Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
-
-            # extract outputs, resize, and remove padding
-            # heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1, 2, 0))  # output 1 is heatmaps
-            heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0))  # output 1 is heatmaps
-            heatmap = util.smart_resize_k(heatmap, fx=stride, fy=stride)
-            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
-            heatmap = util.smart_resize(heatmap, (oriImg.shape[0], oriImg.shape[1]))
-
-            # paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0))  # output 0 is PAFs
-            paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0))  # output 0 is PAFs
-            paf = util.smart_resize_k(paf, fx=stride, fy=stride)
-            paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
-            paf = util.smart_resize(paf, (oriImg.shape[0], oriImg.shape[1]))
-
-            heatmap_avg += heatmap_avg + heatmap / len(multiplier)
-            paf_avg += + paf / len(multiplier)
-
-        all_peaks = []
-        peak_counter = 0
-
-        for part in range(18):
-            map_ori = heatmap_avg[:, :, part]
-            one_heatmap = gaussian_filter(map_ori, sigma=3)
-
-            map_left = np.zeros(one_heatmap.shape)
-            map_left[1:, :] = one_heatmap[:-1, :]
-            map_right = np.zeros(one_heatmap.shape)
-            map_right[:-1, :] = one_heatmap[1:, :]
-            map_up = np.zeros(one_heatmap.shape)
-            map_up[:, 1:] = one_heatmap[:, :-1]
-            map_down = np.zeros(one_heatmap.shape)
-            map_down[:, :-1] = one_heatmap[:, 1:]
-
-            peaks_binary = np.logical_and.reduce(
-                (one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down, one_heatmap > thre1))
-            peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]))  # note reverse
-            peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
-            peak_id = range(peak_counter, peak_counter + len(peaks))
-            peaks_with_score_and_id = [peaks_with_score[i] + (peak_id[i],) for i in range(len(peak_id))]
-
-            all_peaks.append(peaks_with_score_and_id)
-            peak_counter += len(peaks)
-
-        # find connection in the specified sequence, center 29 is in the position 15
-        limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
-                   [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
-                   [1, 16], [16, 18], [3, 17], [6, 18]]
-        # the middle joints heatmap correpondence
-        mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], \
-                  [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], \
-                  [55, 56], [37, 38], [45, 46]]
-
-        connection_all = []
-        special_k = []
-        mid_num = 10
-
-        for k in range(len(mapIdx)):
-            score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
-            candA = all_peaks[limbSeq[k][0] - 1]
-            candB = all_peaks[limbSeq[k][1] - 1]
-            nA = len(candA)
-            nB = len(candB)
-            indexA, indexB = limbSeq[k]
-            if (nA != 0 and nB != 0):
-                connection_candidate = []
-                for i in range(nA):
-                    for j in range(nB):
-                        vec = np.subtract(candB[j][:2], candA[i][:2])
-                        norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
-                        norm = max(0.001, norm)
-                        vec = np.divide(vec, norm)
-
-                        startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
-                                            np.linspace(candA[i][1], candB[j][1], num=mid_num)))
-
-                        vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
-                                          for I in range(len(startend))])
-                        vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
-                                          for I in range(len(startend))])
-
-                        score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
-                        score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
-                            0.5 * oriImg.shape[0] / norm - 1, 0)
-                        criterion1 = len(np.nonzero(score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
-                        criterion2 = score_with_dist_prior > 0
-                        if criterion1 and criterion2:
-                            connection_candidate.append(
-                                [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
-
-                connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
-                connection = np.zeros((0, 5))
-                for c in range(len(connection_candidate)):
-                    i, j, s = connection_candidate[c][0:3]
-                    if (i not in connection[:, 3] and j not in connection[:, 4]):
-                        connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
-                        if (len(connection) >= min(nA, nB)):
-                            break
-
-                connection_all.append(connection)
-            else:
-                special_k.append(k)
-                connection_all.append([])
-
-        # last number in each row is the total parts number of that person
-        # the second last number in each row is the score of the overall configuration
-        subset = -1 * np.ones((0, 20))
-        candidate = np.array([item for sublist in all_peaks for item in sublist])
-
-        for k in range(len(mapIdx)):
-            if k not in special_k:
-                partAs = connection_all[k][:, 0]
-                partBs = connection_all[k][:, 1]
-                indexA, indexB = np.array(limbSeq[k]) - 1
-
-                for i in range(len(connection_all[k])):  # = 1:size(temp,1)
-                    found = 0
-                    subset_idx = [-1, -1]
-                    for j in range(len(subset)):  # 1:size(subset,1):
-                        if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
-                            subset_idx[found] = j
-                            found += 1
-
-                    if found == 1:
-                        j = subset_idx[0]
-                        if subset[j][indexB] != partBs[i]:
-                            subset[j][indexB] = partBs[i]
-                            subset[j][-1] += 1
-                            subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
-                    elif found == 2:  # if found 2 and disjoint, merge them
-                        j1, j2 = subset_idx
-                        membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
-                        if len(np.nonzero(membership == 2)[0]) == 0:  # merge
-                            subset[j1][:-2] += (subset[j2][:-2] + 1)
-                            subset[j1][-2:] += subset[j2][-2:]
-                            subset[j1][-2] += connection_all[k][i][2]
-                            subset = np.delete(subset, j2, 0)
-                        else:  # as like found == 1
-                            subset[j1][indexB] = partBs[i]
-                            subset[j1][-1] += 1
-                            subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
-
-                    # if find no partA in the subset, create a new subset
-                    elif not found and k < 17:
-                        row = -1 * np.ones(20)
-                        row[indexA] = partAs[i]
-                        row[indexB] = partBs[i]
-                        row[-1] = 2
-                        row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
-                        subset = np.vstack([subset, row])
-        # delete some rows of subset which has few parts occur
-        deleteIdx = []
-        for i in range(len(subset)):
-            if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
-                deleteIdx.append(i)
-        subset = np.delete(subset, deleteIdx, axis=0)
-
-        # subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts
-        # candidate: x, y, score, id
-        return candidate, subset
-    
-    @staticmethod
-    def format_body_result(candidate: np.ndarray, subset: np.ndarray) -> List[BodyResult]:
-        """
-        Format the body results from the candidate and subset arrays into a list of BodyResult objects.
-        
-        Args:
-            candidate (np.ndarray): An array of candidates containing the x, y coordinates, score, and id
-                for each body part.
-            subset (np.ndarray): An array of subsets containing indices to the candidate array for each
-                person detected. The last two columns of each row hold the total score and total parts
-                of the person.
-
-        Returns:
-            List[BodyResult]: A list of BodyResult objects, where each object represents a person with
-                detected keypoints, total score, and total parts.
-        """
-        return [
-            BodyResult(
-                keypoints=[
-                    Keypoint(
-                        x=candidate[candidate_index][0],
-                        y=candidate[candidate_index][1],
-                        score=candidate[candidate_index][2],
-                        id=candidate[candidate_index][3]
-                    ) if candidate_index != -1 else None
-                    for candidate_index in person[:18].astype(int)
-                ],
-                total_score=person[18],
-                total_parts=person[19]
-            )
-            for person in subset
-        ]

controlnet_aux_local/open_pose/face.py

@@ -1,364 +0,0 @@
-import logging
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch.nn import Conv2d, MaxPool2d, Module, ReLU, init
-from torchvision.transforms import ToPILImage, ToTensor
-
-from . import util
-
-
-class FaceNet(Module):
-    """Model the cascading heatmaps. """
-    def __init__(self):
-        super(FaceNet, self).__init__()
-        # cnn to make feature map
-        self.relu = ReLU()
-        self.max_pooling_2d = MaxPool2d(kernel_size=2, stride=2)
-        self.conv1_1 = Conv2d(in_channels=3, out_channels=64,
-                              kernel_size=3, stride=1, padding=1)
-        self.conv1_2 = Conv2d(
-            in_channels=64, out_channels=64, kernel_size=3, stride=1,
-            padding=1)
-        self.conv2_1 = Conv2d(
-            in_channels=64, out_channels=128, kernel_size=3, stride=1,
-            padding=1)
-        self.conv2_2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=3, stride=1,
-            padding=1)
-        self.conv3_1 = Conv2d(
-            in_channels=128, out_channels=256, kernel_size=3, stride=1,
-            padding=1)
-        self.conv3_2 = Conv2d(
-            in_channels=256, out_channels=256, kernel_size=3, stride=1,
-            padding=1)
-        self.conv3_3 = Conv2d(
-            in_channels=256, out_channels=256, kernel_size=3, stride=1,
-            padding=1)
-        self.conv3_4 = Conv2d(
-            in_channels=256, out_channels=256, kernel_size=3, stride=1,
-            padding=1)
-        self.conv4_1 = Conv2d(
-            in_channels=256, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv4_2 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv4_3 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv4_4 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv5_1 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv5_2 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv5_3_CPM = Conv2d(
-            in_channels=512, out_channels=128, kernel_size=3, stride=1,
-            padding=1)
-
-        # stage1
-        self.conv6_1_CPM = Conv2d(
-            in_channels=128, out_channels=512, kernel_size=1, stride=1,
-            padding=0)
-        self.conv6_2_CPM = Conv2d(
-            in_channels=512, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage2
-        self.Mconv1_stage2 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage2 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage3
-        self.Mconv1_stage3 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage3 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage4
-        self.Mconv1_stage4 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage4 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage5
-        self.Mconv1_stage5 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage5 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage6
-        self.Mconv1_stage6 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage6 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        for m in self.modules():
-            if isinstance(m, Conv2d):
-                init.constant_(m.bias, 0)
-
-    def forward(self, x):
-        """Return a list of heatmaps."""
-        heatmaps = []
-
-        h = self.relu(self.conv1_1(x))
-        h = self.relu(self.conv1_2(h))
-        h = self.max_pooling_2d(h)
-        h = self.relu(self.conv2_1(h))
-        h = self.relu(self.conv2_2(h))
-        h = self.max_pooling_2d(h)
-        h = self.relu(self.conv3_1(h))
-        h = self.relu(self.conv3_2(h))
-        h = self.relu(self.conv3_3(h))
-        h = self.relu(self.conv3_4(h))
-        h = self.max_pooling_2d(h)
-        h = self.relu(self.conv4_1(h))
-        h = self.relu(self.conv4_2(h))
-        h = self.relu(self.conv4_3(h))
-        h = self.relu(self.conv4_4(h))
-        h = self.relu(self.conv5_1(h))
-        h = self.relu(self.conv5_2(h))
-        h = self.relu(self.conv5_3_CPM(h))
-        feature_map = h
-
-        # stage1
-        h = self.relu(self.conv6_1_CPM(h))
-        h = self.conv6_2_CPM(h)
-        heatmaps.append(h)
-
-        # stage2
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage2(h))
-        h = self.relu(self.Mconv2_stage2(h))
-        h = self.relu(self.Mconv3_stage2(h))
-        h = self.relu(self.Mconv4_stage2(h))
-        h = self.relu(self.Mconv5_stage2(h))
-        h = self.relu(self.Mconv6_stage2(h))
-        h = self.Mconv7_stage2(h)
-        heatmaps.append(h)
-
-        # stage3
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage3(h))
-        h = self.relu(self.Mconv2_stage3(h))
-        h = self.relu(self.Mconv3_stage3(h))
-        h = self.relu(self.Mconv4_stage3(h))
-        h = self.relu(self.Mconv5_stage3(h))
-        h = self.relu(self.Mconv6_stage3(h))
-        h = self.Mconv7_stage3(h)
-        heatmaps.append(h)
-
-        # stage4
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage4(h))
-        h = self.relu(self.Mconv2_stage4(h))
-        h = self.relu(self.Mconv3_stage4(h))
-        h = self.relu(self.Mconv4_stage4(h))
-        h = self.relu(self.Mconv5_stage4(h))
-        h = self.relu(self.Mconv6_stage4(h))
-        h = self.Mconv7_stage4(h)
-        heatmaps.append(h)
-
-        # stage5
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage5(h))
-        h = self.relu(self.Mconv2_stage5(h))
-        h = self.relu(self.Mconv3_stage5(h))
-        h = self.relu(self.Mconv4_stage5(h))
-        h = self.relu(self.Mconv5_stage5(h))
-        h = self.relu(self.Mconv6_stage5(h))
-        h = self.Mconv7_stage5(h)
-        heatmaps.append(h)
-
-        # stage6
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage6(h))
-        h = self.relu(self.Mconv2_stage6(h))
-        h = self.relu(self.Mconv3_stage6(h))
-        h = self.relu(self.Mconv4_stage6(h))
-        h = self.relu(self.Mconv5_stage6(h))
-        h = self.relu(self.Mconv6_stage6(h))
-        h = self.Mconv7_stage6(h)
-        heatmaps.append(h)
-
-        return heatmaps
-
-
-LOG = logging.getLogger(__name__)
-TOTEN = ToTensor()
-TOPIL = ToPILImage()
-
-
-params = {
-    'gaussian_sigma': 2.5,
-    'inference_img_size': 736,  # 368, 736, 1312
-    'heatmap_peak_thresh': 0.1,
-    'crop_scale': 1.5,
-    'line_indices': [
-        [0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6],
-        [6, 7], [7, 8], [8, 9], [9, 10], [10, 11], [11, 12], [12, 13],
-        [13, 14], [14, 15], [15, 16],
-        [17, 18], [18, 19], [19, 20], [20, 21],
-        [22, 23], [23, 24], [24, 25], [25, 26],
-        [27, 28], [28, 29], [29, 30],
-        [31, 32], [32, 33], [33, 34], [34, 35],
-        [36, 37], [37, 38], [38, 39], [39, 40], [40, 41], [41, 36],
-        [42, 43], [43, 44], [44, 45], [45, 46], [46, 47], [47, 42],
-        [48, 49], [49, 50], [50, 51], [51, 52], [52, 53], [53, 54],
-        [54, 55], [55, 56], [56, 57], [57, 58], [58, 59], [59, 48],
-        [60, 61], [61, 62], [62, 63], [63, 64], [64, 65], [65, 66],
-        [66, 67], [67, 60]
-    ],
-}
-
-
-class Face(object):
-    """
-    The OpenPose face landmark detector model.
-
-    Args:
-        inference_size: set the size of the inference image size, suggested:
-            368, 736, 1312, default 736
-        gaussian_sigma: blur the heatmaps, default 2.5
-        heatmap_peak_thresh: return landmark if over threshold, default 0.1
-
-    """
-    def __init__(self, face_model_path,
-                 inference_size=None,
-                 gaussian_sigma=None,
-                 heatmap_peak_thresh=None):
-        self.inference_size = inference_size or params["inference_img_size"]
-        self.sigma = gaussian_sigma or params['gaussian_sigma']
-        self.threshold = heatmap_peak_thresh or params["heatmap_peak_thresh"]
-        self.model = FaceNet()
-        self.model.load_state_dict(torch.load(face_model_path))
-        self.model.eval()
-
-    def to(self, device):
-        self.model.to(device)
-        return self
-
-    def __call__(self, face_img):
-        device = next(iter(self.model.parameters())).device
-        H, W, C = face_img.shape
-
-        w_size = 384
-        x_data = torch.from_numpy(util.smart_resize(face_img, (w_size, w_size))).permute([2, 0, 1]) / 256.0 - 0.5
-
-        x_data = x_data.to(device)
-
-        with torch.no_grad():
-            hs = self.model(x_data[None, ...])
-            heatmaps = F.interpolate(
-                hs[-1],
-                (H, W),
-                mode='bilinear', align_corners=True).cpu().numpy()[0]
-        return heatmaps
-
-    def compute_peaks_from_heatmaps(self, heatmaps):
-        all_peaks = []
-        for part in range(heatmaps.shape[0]):
-            map_ori = heatmaps[part].copy()
-            binary = np.ascontiguousarray(map_ori > 0.05, dtype=np.uint8)
-
-            if np.sum(binary) == 0:
-                continue
-
-            positions = np.where(binary > 0.5)
-            intensities = map_ori[positions]
-            mi = np.argmax(intensities)
-            y, x = positions[0][mi], positions[1][mi]
-            all_peaks.append([x, y])
-
-        return np.array(all_peaks)