Create app.py

app.py

@@ -0,0 +1,542 @@
+import os
+os.system('pip install numpy pillow opencv-python fastapi starlette uvicorn requests')
+import gradio as gr
+import shutil
+import cv2
+import numpy as np
+from PIL import Image
+import asyncio
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from gradio_client import Client
+from fastapi import FastAPI, Request
+from fastapi.staticfiles import StaticFiles
+from starlette.responses import Response
+import uvicorn
+import requests
+from urllib.parse import quote
+from unicodedata import normalize
+
+root = os.path.dirname(os.path.abspath(__file__))
+textures_folder = os.path.join(root, 'textures')
+os.makedirs(textures_folder, exist_ok=True)
+valid_extensions = ['.jpeg', '.jpg', '.png']
+
+
+
+textures_repo = "https://huggingface.co/datasets/2ch/textures/resolve/main/"
+textures_for_download = [
+f"{textures_repo}гауссовский_шум_и_мелкое_зерно.png?download=true",
+f"{textures_repo}грязная_матрица.png?download=true",
+f"{textures_repo}для_ночных_и_тёмных_кадров_сильный_шум_и_пыль.png?download=true",
+f"{textures_repo}для_ночных_и_тёмных_кадров_царапины_шум_пыль_дымка.png?download=true",
+f"{textures_repo}для_светлых_и_солнечных_ярких_фото_мелкое_констрастное_зерно.png?download=true",
+f"{textures_repo}зернистость_плёнки.png?download=true",
+f"{textures_repo}зернистость_плёнки_с_грязью.png?download=true",
+f"{textures_repo}испорченная_ворсом_плёнка.png?download=true",
+f"{textures_repo}мелкий_цветной_шум.png?download=true",
+f"{textures_repo}мелкое_контрастное_зерно_и_средний_цветвой_шум.png?download=true",
+f"{textures_repo}очень_мелкое_зерно.png?download=true",
+f"{textures_repo}пыльная_плёнка.png?download=true",
+f"{textures_repo}сильный_цветовой_шум_для_ночных_фото.png?download=true",
+f"{textures_repo}слабый_естественный_шум_матрицы_смартфона.png?download=true",
+f"{textures_repo}среднее_зерно.png?download=true",
+f"{textures_repo}среднее_монохромное_зерно_пыль_и_ворсинки.png?download=true",
+f"{textures_repo}средний_цветной_шум.png?download=true",
+f"{textures_repo}старая_матрица.png?download=true",
+f"{textures_repo}старая_потёртая_плёнка.png?download=true",
+f"{textures_repo}цветной_шум_матрицы.png?download=true",
+f"{textures_repo}цветной_шум_на_плёнке.png?download=true",
+f"{textures_repo}шумная_матрица.png?download=true",
+    ]
+
+
+def dl_textures(texture_url):
+    texture_for_download = quote(normalize('NFD', texture_url), safe='/?:=')
+    filename = texture_url.split('/')[-1].split('?')[0]
+    file_path = os.path.join(textures_folder, filename)
+    response = requests.get(texture_for_download, stream=True)
+    response.raise_for_status()
+    with open(file_path, 'wb') as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+
+
+def create_texture_preview(texture_folder, output_folder, size=(246, 246)):
+    os.makedirs(output_folder, exist_ok=True)
+    for texture in os.listdir(texture_folder):
+        img_path = os.path.join(texture_folder, texture)
+        img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+        start_x = np.random.randint(0, img.shape[1] - size[1])
+        start_y = np.random.randint(0, img.shape[0] - size[0])
+        img = img[start_y:start_y + size[0], start_x:start_x + size[1]]
+        cv2.imwrite(os.path.join(output_folder, texture), img)
+
+
+def prepare_textures(texture_folder, output_folder):
+    with ThreadPoolExecutor(max_workers=len(textures_for_download)) as executor:
+        futures = [executor.submit(dl_textures, texture_for_download) for texture_for_download in
+                   textures_for_download]
+        for future in as_completed(futures):
+            future.result()
+
+    create_texture_preview(texture_folder, output_folder, size=(246, 246))
+
+
+prepare_textures(textures_folder, os.path.join(root, 'preview'))
+
+
+
+preview_css = ""
+for i, texture in enumerate(os.listdir(textures_folder), start=1):
+    if os.path.splitext(texture)[1].lower() in valid_extensions:
+        preview_css += f"""[data-testid="{i:02d}-radio-label"]::before {{
+          background-color: transparent !important;
+          background-image: url("./preview/{texture}") !important;
+        }}\n"""
+
+radio_css = """
+html,
+body {
+    background: var(--body-background-fill);
+}
+
+.gradio-container {
+    max-width: 1396px !important;
+}
+
+#textures label {
+    position: relative;
+    width: 256px;
+    height: 256px;
+    display: flex;
+    flex-direction: row;
+    align-items: flex-end;
+    background: none !important;
+    padding: 4px !important;
+    transition: .3s;
+}
+
+#textures label::before {
+    width: 246px;
+    height: 246px;
+    border-radius: 8px;
+    display: block;
+    content: "";
+    transition: .3s;
+    background: red;
+    position: relative;
+    top: 0px;
+}
+
+#textures label:hover::before,
+#textures label:active::before,
+#textures label.selected::before {
+    mix-blend-mode: soft-light;
+    transition: .3s
+}
+
+#textures span:not([data-testid="block-info"]),
+#textures input {
+    position: absolute;
+    z-index: 999;
+}
+
+#textures input {
+    position: absolute;
+    z-index: 999;
+    bottom: 9px;
+    left: 9px;
+}
+
+#textures span:not([data-testid="block-info"]) {
+    left: 21px;
+    padding: 2px 8px;
+    background: rgba(0, 0, 0, .57);
+    backdrop-filter: blur(3px)
+}
+
+#textures {
+    background-color: hsla(0, 0%, 50%, 1);
+}
+
+.built-with,
+.show-api,
+footer .svelte-mpyp5e {
+    display: none !important;
+}
+
+footer:after {
+    content: "ну пролапс, ну и что?";
+}
+
+#zoom {
+    position: absolute;
+    top: 50%;
+    left: 50%;
+    width: 250px;
+    height: 250px;
+    background-repeat: no-repeat;
+    box-shadow: 0px 0px 10px 5px rgba(0, 0, 0, .2);
+    border-radius: 50%;
+    cursor: none;
+    pointer-events: none;
+    z-index: 999;
+    opacity: 0;
+    transform: scale(0);
+    transition: opacity 500ms, transform 500ms;
+}
+
+#textures_tab .image-button {
+    cursor: none;
+}
+
+#textured_result-download-link,
+#restored_image-download-link {
+    position: absolute;
+    z-index: 9999;
+    padding: 2px 4px;
+    margin: 0 7px;
+    background: black;
+    bottom: 0;
+    right: 0;
+    font-size: 20px;
+    transition: 300ms
+}
+
+#download-link:hover {
+    color: #99f7a8
+}
+
+#restored_images.disabled {
+    height: 0px !important;
+    opacity: 0;
+    transition: 300ms
+}
+
+#restored_images.enabled {
+    transition: 300ms
+}
+""" + preview_css
+
+custom_js = """
+const PageLoadObserver = new MutationObserver((mutationsList, observer) => {
+    for (let mutation of mutationsList) {
+        if (mutation.type === 'childList') {
+            const tabsDiv = document.querySelector('div.tab-nav');
+            if (tabsDiv) {
+                observer.disconnect();
+                document.querySelector('#textures_tab-button').addEventListener('click', () => {
+                    setTimeout(() => {
+                        let labels = document.querySelectorAll('label[data-testid]');
+                        labels.forEach((label) => {
+                            let input = label.querySelector('input[type="radio"]');
+                            if (input) {
+                                let title = input.value.split('.')[0].replace(/_/g, ' ');
+                                label.title = title;
+                            }
+                        });
+                        document.querySelector("label[data-testid='05-radio-label']").click()
+                    }, 150);
+                })
+                let RestoredGallery = document.getElementById('restored_images');
+                function checkImagesAndSetClass() {
+                    const firstDiv = RestoredGallery.querySelector('div:first-child');
+                    const hasChildElements = firstDiv && firstDiv.children.length > 0;
+                    const hasImages = RestoredGallery.querySelectorAll('img').length > 0;
+                    if (hasChildElements || hasImages) {
+                        RestoredGallery.classList.add('enabled');
+                        RestoredGallery.classList.remove('disabled');
+                    } else {
+                        RestoredGallery.classList.add('disabled');
+                        RestoredGallery.classList.remove('enabled');
+                    }
+                }
+                const FaceResoreResultCheck = new MutationObserver((mutations) => {
+                checkImagesAndSetClass();
+                });
+                FaceResoreResultCheck.observe(RestoredGallery, {childList: true, subtree: true});
+                checkImagesAndSetClass();
+                function magnify(imgID, zoom) {
+                    var img, glass, w, h, bw;
+                    img = document.querySelector(imgID);
+                    glass = document.createElement("DIV");
+                    glass.setAttribute("id", "zoom");
+                    img.parentElement.insertBefore(glass, img);
+                    glass.style.backgroundImage = "url('" + img.src + "')";
+                    glass.style.backgroundRepeat = "no-repeat";
+                    glass.style.backgroundSize = (img.width * zoom) + "px " + (img.height * zoom) + "px";
+                    bw = 3;
+                    w = glass.offsetWidth / 2;
+                    h = glass.offsetHeight / 2;
+                    glass.addEventListener("mousemove", moveMagnifier);
+                    img.addEventListener("mousemove", moveMagnifier);
+                    glass.addEventListener("touchmove", moveMagnifier);
+                    img.addEventListener("touchmove", moveMagnifier);
+                    function moveMagnifier(e) {
+                        var pos, x, y;
+                        e.preventDefault();
+                        pos = getCursorPos(e);
+                        x = pos.x;
+                        y = pos.y;
+                        if (x > img.width - (w / zoom)) { x = img.width - (w / zoom); }
+                        if (x < w / zoom) { x = w / zoom; }
+                        if (y > img.height - (h / zoom)) { y = img.height - (h / zoom); }
+                        if (y < h / zoom) { y = h / zoom; }
+                        glass.style.left = (x - w) + "px";
+                        glass.style.top = (y - h) + "px";
+                        glass.style.backgroundPosition = "-" + ((x * zoom) - w + bw) + "px -" + ((y * zoom) - h) + "px";
+                        glass.style.backgroundImage = "url('" + img.src + "')";
+                    }
+                    function getCursorPos(e) {
+                        var a, x = 0, y = 0;
+                        e = e || window.event;
+                        a = img.getBoundingClientRect();
+                        x = e.pageX - a.left;
+                        y = e.pageY - a.top;
+                        x = x - window.scrollX;
+                        y = y - window.scrollY;
+                        return { x: x, y: y };
+                    }
+
+                    img.addEventListener("mouseover", function () {
+
+                        glass.style.opacity = "1";
+                        glass.style.transform = "scale(1)";
+                    });
+                    img.addEventListener("mouseout", function () {
+                        glass.style.opacity = "0";
+                        glass.style.transform = "scale(0)";
+                    });
+                }
+                function setupDownloadLink(imgSelector, linkSelector, linkId, magnifyImage) {
+                    const imgElement = document.querySelector(imgSelector);
+                    if (imgElement && imgElement.src) {
+                      let downloadLink = document.querySelector(linkSelector);
+                      if (!downloadLink) {
+                        if (magnifyImage) {
+                          magnify(magnifyImage, 3);
+                        }
+                        downloadLink = document.createElement('a');
+                        downloadLink.id = linkId;
+                        downloadLink.innerText = 'скачать';
+                        imgElement.after(downloadLink);
+                      }
+                      downloadLink.href = imgElement.src;
+                      downloadLink.download = '';
+                    }
+                  }
+
+                  const DownloadLinkObserverCallback = (mutationsList, observer, imgSelector, linkSelector, linkId, magnifyImage) => {
+                    setupDownloadLink(imgSelector, linkSelector, linkId, magnifyImage);
+                  };
+
+                  const DownloadLinkObserverOptions = { childList: true, subtree: true, attributes: true, attributeFilter: ['src'] };
+
+                  const ImageTexturedObserver = new MutationObserver((mutationsList, observer) => {
+                    DownloadLinkObserverCallback(mutationsList, observer, '#textured_result img[data-testid="detailed-image"]', '#textured_result-download-link', 'textured_result-download-link', "#textured_result .image-button img");
+                  });
+
+                  ImageTexturedObserver.observe(document, DownloadLinkObserverOptions);
+
+                  const ImageRestoredObserver = new MutationObserver((mutationsList, observer) => {
+                    DownloadLinkObserverCallback(mutationsList, observer, '#restored_images img[data-testid="detailed-image"]', '#restored_image-download-link', 'restored_image-download-link');
+                  });
+
+                  ImageRestoredObserver.observe(document, DownloadLinkObserverOptions);
+
+            }
+        }
+    }
+});
+
+PageLoadObserver.observe(document, { childList: true, subtree: true });
+"""
+
+def extract_path_from_result(predict_answer):
+    if isinstance(predict_answer, (tuple, list)):
+        result = predict_answer[0]
+        shutil.rmtree(os.path.dirname(predict_answer[1]), ignore_errors=True)
+    else:
+        result = predict_answer
+    return result
+
+
+def restore_face_common(img_path: str, predict_answer: str, model: str) -> None:
+    result = extract_path_from_result(predict_answer)
+    if os.path.exists(result):
+        if os.path.exists(img_path):
+            os.unlink(img_path)
+        new_file, new_extension = os.path.splitext(result)
+        old_file, old_extension = os.path.splitext(img_path)
+        old_filename = os.path.basename(old_file)
+        new_location = os.path.join(os.path.dirname(img_path), f"{old_filename}_{model}{new_extension}")
+        shutil.move(result, new_location)
+        shutil.rmtree(os.path.dirname(result), ignore_errors=True)
+
+
+def restore_face_gfpgan(img_path: str) -> None:
+    client = Client(src="https://xintao-gfpgan.hf.space/", verbose=False)
+    result = client.predict(img_path, "v1.4", 4, api_name="/predict")
+    restore_face_common(img_path, result, "gfpgan")
+
+
+def restore_face_codeformer(img_path: str) -> None:
+    client = Client(src="https://sczhou-codeformer.hf.space/", verbose=False)
+    result = client.predict(img_path, True, True, True, 2, 0, api_name="/predict")
+    restore_face_common(img_path, result, "codeformer")
+
+
+
+async def restore_faces_one_image(img_path: str, func_list: list) -> bool:
+    def run_func(func) -> bool:
+        for _ in range(3):
+            try:
+                func(img_path)
+                return True
+            except Exception as e:
+                print(f"ошибка в {func.__name__}: {e}")
+        return False
+
+    loop = asyncio.get_event_loop()
+    with ThreadPoolExecutor(max_workers=len(func_list)) as executor:
+        futures = [loop.run_in_executor(executor, run_func, func) for func in func_list]
+    results = await asyncio.gather(*futures)
+    return any(results)
+
+
+async def restore_faces_batch(input_images: list[str], func_list: list, batch_size: int = 3) -> bool:
+    results = False
+    try:
+        batches = [input_images[i:i + batch_size] for i in range(0, len(input_images), batch_size)]
+        for batch in batches:
+            tasks = [restore_faces_one_image(img_path, func_list) for img_path in batch]
+            results = await asyncio.gather(*tasks)
+        return any(results)
+    except Exception as error:
+        print(error)
+        return results
+
+
+def get_file_paths(input_path: str | list[str], extensions_list: list[str]) -> list[str]:
+    files = []
+
+    def add_files_from_directory(directory):
+        for file_name in os.listdir(directory):
+            if os.path.splitext(file_name)[1] in extensions_list:
+                files.append(os.path.abspath(os.path.join(directory, file_name)))
+
+    if isinstance(input_path, list):
+        for file_path in input_path:
+            parent_directory = os.path.dirname(file_path)
+            add_files_from_directory(parent_directory)
+    else:
+        add_files_from_directory(input_path)
+
+    return files
+
+
+async def restore_upscale(files, restore_method):
+    file_paths = [file.name for file in files]
+    if restore_method == 'codeformer':
+        func_list = [restore_face_codeformer]
+    elif restore_method == 'gfpgan':
+        func_list = [restore_face_gfpgan]
+    else:
+        func_list = [restore_face_codeformer, restore_face_gfpgan]
+    results = await restore_faces_batch(file_paths, func_list, batch_size=3)
+    if results:
+        file_paths = get_file_paths(file_paths, valid_extensions)
+        print(f"restore_upscale: get_file_paths: {file_paths}")
+        return file_paths
+    else:
+        return [os.path.join(root, 'error.png')]
+
+
+def image_noise_softlight_layer_mix(img, texture, output: str = None, opacity: float = 0.7):
+    if isinstance(img, Image.Image):
+        img = np.array(img).astype(float)
+    elif isinstance(img, np.ndarray):
+        img = img.astype(float)
+
+    if img.shape[2] == 3 and not isinstance(img, Image.Image):
+        img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_RGB2BGR).astype(float)
+
+    overlay = cv2.imread(texture, cv2.IMREAD_UNCHANGED).astype(float)
+    start_x = np.random.randint(0, overlay.shape[1] - img.shape[1])
+    start_y = np.random.randint(0, overlay.shape[0] - img.shape[0])
+    overlay = overlay[start_y:start_y + img.shape[0], start_x:start_x + img.shape[1]]
+    if img.shape[2] == 3:
+        img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_RGB2RGBA).astype(float)
+    if overlay.shape[2] == 3:
+        overlay = cv2.cvtColor(overlay.astype(np.uint8), cv2.COLOR_RGB2RGBA).astype(float)
+    overlay[..., 3] *= opacity
+    img_in_norm = img / 255.0
+    img_layer_norm = overlay / 255.0
+    comp_alpha = np.minimum(img_in_norm[:, :, 3], img_layer_norm[:, :, 3]) * 1.0
+    new_alpha = img_in_norm[:, :, 3] + (1.0 - img_in_norm[:, :, 3]) * comp_alpha
+    np.seterr(divide='ignore', invalid='ignore')
+    ratio = comp_alpha / new_alpha
+    ratio[ratio == np.NAN] = 0.0
+    comp = (1.0 - img_in_norm[:, :, :3]) * img_in_norm[:, :, :3] * img_layer_norm[:, :, :3] + img_in_norm[:, :, :3] * (
+            1.0 - (1.0 - img_in_norm[:, :, :3]) * (1.0 - img_layer_norm[:, :, :3]))
+    ratio_rs = np.reshape(np.repeat(ratio, 3), [comp.shape[0], comp.shape[1], comp.shape[2]])
+    img_out = comp * ratio_rs + img_in_norm[:, :, :3] * (1.0 - ratio_rs)
+    img_out = np.nan_to_num(np.dstack((img_out, img_in_norm[:, :, 3])))
+    result = img_out * 255.0
+    rgb_image = cv2.cvtColor(result.astype(np.uint8), cv2.COLOR_BGR2RGB)
+    image = Image.fromarray(rgb_image)
+    return np.array(image)
+
+
+def apply_texture(input_image, textures_choice, opacity_slider):
+    result = image_noise_softlight_layer_mix(input_image, os.path.join(textures_folder, textures_choice), opacity=opacity_slider)
+    return [result]
+
+
+with gr.Blocks(analytics_enabled=False, css=radio_css) as demo:
+    with gr.Tab(label="восстановление лиц", id=1, elem_id="restore_tab"):
+        restore_method = gr.Radio(["codeformer", "gfpgan", "оба"], value="codeformer", label="", interactive=True)
+        restore_method.change(fn=lambda x: print(f"restore_method value = {x}"), inputs=restore_method, api_name="show_selected_method")
+        file_output = gr.Gallery(label="", container=True, object_fit="cover", columns=4, rows=4, allow_preview=True, preview=True, show_share_button=False, show_download_button=False, elem_id="restored_images")
+        upload_button = gr.UploadButton("выбор изображений для обработки", file_types=["image"], file_count="multiple", variant="primary")
+        upload_button.upload(fn=restore_upscale, inputs=[upload_button, restore_method], outputs=file_output, api_name="face_restore")
+    with gr.Tab(label="наложение зернистости пленки и шума", id=2, elem_id="textures_tab"):
+        with gr.Row(variant="compact", elem_id="textures_tab_images"):
+            input_image = gr.Image(label="исходник", sources=["upload", "clipboard"], type="numpy")
+            result_image = gr.Gallery(label="результат", elem_id="textured_result", allow_preview=True, preview=True, show_share_button=False, show_download_button=False)
+        opacity_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="видимость")
+        apply_button = gr.Button(value="применить", variant="primary")
+        texture_files = [(f"{i:02d}", texture) for i, texture in enumerate(os.listdir(textures_folder), start=1) if os.path.splitext(texture)[1].lower() in valid_extensions]
+        textures_choice = gr.Radio(texture_files, show_label=False, interactive=True, elem_id="textures")
+        apply_button.click(fn=apply_texture, inputs=[input_image, textures_choice, opacity_slider], outputs=result_image, api_name="texturize")
+
+
+app = FastAPI()
+
+
+@app.middleware("http")
+async def some_fastapi_middleware(request: Request, call_next):
+    response = await call_next(request)
+    path = request.url.path
+    if path == "/":
+        response_body = ""
+        async for chunk in response.body_iterator:
+            response_body += chunk.decode()
+        javascript = f"""
+        <script type="text/javascript">
+        {custom_js}
+        </script>
+        """
+        response_body = response_body.replace("</body>", javascript + "</body>")
+        del response.headers["content-length"]
+        return Response(
+            content=response_body,
+            status_code=response.status_code,
+            headers=dict(response.headers),
+            media_type=response.media_type
+        )
+    return response
+
+app.mount("/preview", StaticFiles(directory=os.path.join(root, 'preview')), name="preview")
+gr.mount_gradio_app(app, demo, path="/")
+uvicorn.run(app, host="0.0.0.0", port=7860)