Initial Commit

.gitignore

@@ -0,0 +1,13 @@
+ossutil_output
+__pycache__
+.DS_Store
+test.jpg
+test.py
+test.sh
+app.sh
+__pycache__
+r1.jpg
+r2.png
+r11.png
+test.jpg
+temp.py

app.py

@@ -0,0 +1,398 @@
+import hashlib
+import time
+import uuid
+from urllib.parse import urlencode
+import json
+import requests
+from PIL import Image, ImageOps
+import os
+import gradio as gr
+
+import ops
+from watermark import WatermarkApp
+import db_examples
+
+
+class ApiClient(object):
+
+    def __init__(self, app_key: str, access_key_id: str, access_key_secret: str, endpoint: str):
+        self.app_key = app_key
+        self.access_key_id = access_key_id
+        self.access_key_secret = access_key_secret
+        self.endpoint = endpoint
+        self.base_url = 'http://' + self.endpoint + '/api'
+        self.timeout = 8000
+        self.session = requests.Session()
+        self.session.headers.update(
+            {
+                "Content-Type": "application/json;charset=utf-8",
+                "accessKey": self.access_key_id,
+                "appKey": self.app_key
+            }
+        )
+
+    def send_get(self, headers=None, params=None):
+        if headers is None:
+            headers = {}
+        if params is None:
+            params = {}
+        args = self.cleanNoneValue(
+            {
+                "url": self.base_url,
+                "headers": self._prepare_headers(headers),
+                "params": self._prepare_params(params),
+                "timeout": self.timeout
+            }
+        )
+        response = self._dispatch_request("GET")(**args)
+        self._handle_exception(response)
+        try:
+            data = response.json()
+        except ValueError:
+            data = response.text
+        return data
+
+    def send_post(self, headers=None, params=None):
+        if headers is None:
+            headers = {}
+        if params is None:
+            params = {}
+        args = self.cleanNoneValue(
+            {
+                "url": self.base_url,
+                "headers": self._prepare_headers(headers),
+                "json": params,
+                "timeout": self.timeout
+            }
+        )
+        response = self._dispatch_request("POST")(**args)
+        self._handle_exception(response)
+        try:
+            data = response.json()
+        except ValueError:
+            data = response.text
+        return data
+
+    def _prepare_headers(self, headers):
+        headers['requestId'] = str(uuid.uuid1())
+        timestamp = int(round(time.time() * 1000))
+        headers['timestamp'] = str(timestamp)
+        headers['sign'] = self._get_sign(timestamp)
+        return headers
+
+    def _prepare_params(self, params):
+        return self.encoded_string(self.cleanNoneValue(params))
+
+    def _get_sign(self, timestamp):
+        key = self.app_key + self.access_key_id + str(timestamp) + self.access_key_secret
+        md5hash = hashlib.md5(str.encode(key, 'utf-8'))
+        return md5hash.hexdigest()
+
+    def _dispatch_request(self, http_method):
+        return {
+            "GET": self.session.get,
+            "DELETE": self.session.delete,
+            "PUT": self.session.put,
+            "POST": self.session.post,
+        }.get(http_method, "GET")
+
+    def _handle_exception(self, response):
+        status_code = response.status_code
+        if status_code < 400:
+            return
+            raise Exception(response.text)
+
+    def encoded_string(self, query):
+
+        return urlencode(query, True).replace("%40", "@")
+
+    def cleanNoneValue(self, d) -> dict:
+        out = {}
+        for k in d.keys():
+            if d[k] is not None:
+                out[k] = d[k]
+        return out
+    
+
+api_key = os.environ['APIKEY']
+ak = os.environ['AK']
+sk = os.environ['SK']
+endpoint = os.environ['ENDPOINT']
+apiname = os.environ['APINAME']
+callbackapiname = os.environ['CALLBACKAPINAME']
+osssigneapiname = os.environ['OSSSIGNAPINAME']
+
+client = ApiClient(api_key, ak, sk, endpoint)
+watermark_app = WatermarkApp()
+
+def upload_to_oss(file_path, accessid, policy, signature, host, key, callback):
+    with open(file_path, 'rb') as f:
+        files = {'file': (key, f)}
+        data = {
+            'OSSAccessKeyId': accessid,
+            'policy': policy,
+            'Signature': signature,
+            'key': key,
+            'callback': callback,
+        }
+        response = requests.post(host, files=files, data=data)
+        if response.status_code == 204 or response.ok:
+            return key
+        else:
+            print(f"file upload failed: {response.text}")
+            return None
+
+
+def upload_oss_bucket(image_pil):
+    headers = {
+        'apiName': osssigneapiname
+    }
+    params = {
+        'fileType': '1'
+    }
+    result = client.send_get(headers, params)
+    try:
+        result = result['data']
+    except Exception as e:
+        raise ValueError('oss sign error')
+    accessid = result['accessid']
+    policy = result['policy']
+    signature = result['signature']
+    host = result['host']
+    callback = ''
+    oss_key = os.path.join(result['dir'], '{}.jpg'.format(result['expire']))
+    file_path = 'test.jpg'
+    image_pil.save(file_path)
+    oss_key = upload_to_oss(file_path, accessid, policy, signature, host, oss_key, callback)
+    if oss_key is None:
+        raise ValueError('oss upload error')
+    return oss_key
+    
+
+def call_text_guided_relighting(image, mode, prompt, seed, steps):
+    headers = {
+        'apiName': apiname
+    }
+    image = ImageOps.exif_transpose(image).convert('RGB')
+    image = ops.resize_keep_hw_rate(image, tar_res=1280)
+    image = upload_oss_bucket(image)
+    ops.print_with_datetime('start call_text_guided_relighting')
+    ops.print_with_datetime(prompt)
+    params = {
+        'image': image,
+        'inference_mode': 'free_txt2bg_gen',
+        'mode': mode,
+        'prompt': prompt,
+        'seed': seed,
+        'steps': steps,
+    }
+    ops.print_with_datetime(f'length params, {len(json.dumps(params))}')
+    task_id = client.send_post(headers, params)['data']
+    time.sleep(10)
+    headers = {
+        'apiName': callbackapiname
+    }
+    params = {
+        'id': task_id
+    }
+    flag = True
+    while flag:
+        result = client.send_get(headers, params)['data']
+        if result['status'] != 1:
+            flag = False
+        else:
+            time.sleep(10)
+    if result['status'] != 2:
+        raise ValueError('something wrong in the process')
+    
+    result_1 = result['sasMyCreationPicVOs'][0]['picUrl']
+    result_2 = result['sasMyCreationPicVOs'][0]['maskUrl']
+    result_1 = ops.decode_img_from_url(result_1)
+    result_1 = watermark_app.process_image(result_1)
+    result_2 = ops.decode_img_from_url(result_2)
+    return result_1, result_2
+
+
+def call_image_guided_relighting(image, ref_img, seed, steps):
+    headers = {
+        'apiName': apiname
+    }
+    image = ImageOps.exif_transpose(image).convert('RGB')
+    image = ops.resize_keep_hw_rate(image, tar_res=1280)
+    image = upload_oss_bucket(image)
+    
+    ref_img = ImageOps.exif_transpose(ref_img).convert('RGB')
+    ref_img = ops.resize_keep_hw_rate(ref_img, tar_res=1280)
+    ref_img = upload_oss_bucket(ref_img)
+    
+    ops.print_with_datetime('start call_image_guided_relighting')
+    params = {
+        'image': image,
+        'inference_mode': 'replica_gen',
+        'mode': 'normal',
+        'ref_img': ref_img,
+        'seed': seed,
+        'steps': steps,
+    }
+    ops.print_with_datetime(f'length params, {len(json.dumps(params))}')
+    task_id = client.send_post(headers, params)
+    print(task_id)
+    task_id = task_id['data']
+    time.sleep(10)
+    headers = {
+        'apiName': callbackapiname
+    }
+    params = {
+        'id': task_id
+    }
+    flag = True
+    while flag:
+        result = client.send_get(headers, params)
+        result = result['data']
+        if result['status'] != 1:
+            flag = False
+        else:
+            time.sleep(10)
+    if result['status'] != 2:
+        raise ValueError('something wrong in the process')
+    
+    result_1 = result['sasMyCreationPicVOs'][0]['picUrl']
+    result_2 = result['sasMyCreationPicVOs'][0]['maskUrl']
+    result_1 = ops.decode_img_from_url(result_1)
+    result_1 = watermark_app.process_image(result_1)
+    result_2 = ops.decode_img_from_url(result_2)
+    return result_1, result_2
+
+
+
+quick_prompts = [
+    'warm lighting',
+    'sunshine from window',
+    'neon lighting',
+    'at noon. bright sunlight',
+    'at dusk',
+    'golden time',
+    'natural lighting',
+    'shadow from window',
+    'soft studio lighting',
+    'red lighting',
+    'purple lighting'
+]
+quick_prompts = [[x] for x in quick_prompts]
+
+
+quick_content_prompts = [
+    'by the sea',
+    'in the forest',
+    'on the snow mountain',
+    'by the city street',
+    'on the grassy field',
+    'cityscape',
+    'on the desert',
+    'in the living room',
+]
+quick_content_prompts = [[x] for x in quick_content_prompts]
+
+
+quick_subjects = [
+    'portrait photography of a woman',
+    'portrait photography of man',
+    'product photography',
+]
+quick_subjects = [[x] for x in quick_subjects]
+
+
+with gr.Blocks().queue() as demo:
+    gr.HTML("""
+               <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
+                   FreeLighting: relighting model with both text-condition and image-condition
+               </div>
+               """)
+    with gr.Row():
+        gr.Markdown("See more information in https://github.com/liuyuxuan3060/FreeLighting")
+    with gr.Row():
+        gr.Markdown("We use a open source segmentation model to generate image mask")
+
+    with gr.Tabs():
+        with gr.TabItem("text-guided relighting") as t2v_tab:
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        image = gr.Image(label="original_image", type="pil", height=480)
+                        image_mask = gr.Image(label="image_mask", type="pil", height=480)
+                    with gr.Row():
+                        prompt = gr.Textbox(value="", label="text prompt")     
+                    with gr.Row():
+                        mode = gr.Radio(choices=["normal", "uniform-lit"], value='normal', label="uniform-lit mode will use double time", type='value')
+                        seed = gr.Number(value=12345, label="random seed", precision=0)
+                        steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        button = gr.Button("generate")
+                    
+                    with gr.Row():
+                        example_quick_subjects = gr.Dataset(samples=quick_subjects, label='Subject Quick List', samples_per_page=1000, components=[prompt])
+                    with gr.Row():
+                        example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Lighting Quick List', samples_per_page=1000, components=[prompt])
+                    with gr.Row():
+                        example_quick_content_prompts = gr.Dataset(samples=quick_content_prompts, label='Content Quick List', samples_per_page=1000, components=[prompt])               
+                with gr.Column():
+                    relighting_image = gr.Image(label="relighted_image", type="pil", height=480)
+       
+            with gr.Row():
+                gr.Examples(
+                    examples=db_examples.text_guided_examples,
+                    inputs=[
+                        image, mode, prompt, seed, steps
+                        , relighting_image
+                    ],
+                    outputs=[relighting_image],
+                    examples_per_page=1024
+                )
+                
+            button.click(
+                fn=call_text_guided_relighting,
+                inputs=[
+                    image, mode, prompt, seed, steps
+                ],
+                outputs=[relighting_image, image_mask],
+            )     
+            example_quick_content_prompts.click(lambda x, y: ', '.join(y.split(', ')[:2] + [x[0]]), inputs=[example_quick_content_prompts, prompt], outputs=prompt, show_progress=False, queue=False)
+            example_quick_prompts.click(lambda x, y: ', '.join(y.split(', ')[:2] + [x[0]]), inputs=[example_quick_prompts, prompt], outputs=prompt, show_progress=False, queue=False)
+            example_quick_subjects.click(lambda x: x[0], inputs=example_quick_subjects, outputs=prompt, show_progress=False, queue=False)
+                    
+        with gr.TabItem("image-guided relighting") as i2v_tab:
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        image = gr.Image(label="original_image", type="pil", height=480)
+                        ref_img = gr.Image(label="reference_image", type="pil", height=480)
+                        image_mask = gr.Image(label="image_mask", type="pil", height=480)
+                    with gr.Row():
+                        seed = gr.Number(value=12345, label="random seed", precision=0)
+                        steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
+                        button = gr.Button("generate")
+                with gr.Column():
+                    relighting_image = gr.Image(label="relighted_image", type="pil", height=480)
+
+            with gr.Row():
+                gr.Examples(
+                    examples=db_examples.image_guided_examples,
+                    inputs=[
+                        image, ref_img, seed, steps
+                        , relighting_image
+                    ],
+                    outputs=[relighting_image],
+                    examples_per_page=1024
+                )
+
+            button.click(
+                fn=call_image_guided_relighting,
+                inputs=[
+                    image, ref_img, seed, steps
+                ],
+                outputs=[relighting_image, image_mask],
+            )   
+
+
+demo.launch()
+
+

db_examples.py

@@ -0,0 +1,98 @@
+text_guided_examples = [
+    [
+        'images/original_1.jpg',
+        'uniform-lit', 
+        'portrait photography of a woman, at noon. bright sunlight, by the sea', 
+        1, 
+        20,
+        'images/restxt_1.webp',
+    ],
+    [
+        'images/original_1.jpg',
+        'uniform-lit', 
+        'portrait photography of a woman, neon lighting, by the city street', 
+        1, 
+        20,
+        'images/restxt_2.webp',
+    ],
+    [
+        'images/original_2.jpg',
+        'normal', 
+        'portrait photography of a woman, purple lighting, in the living room', 
+        1, 
+        20,
+        'images/restxt_3.webp',
+    ],
+    [
+        'images/original_2.jpg',
+        'normal', 
+        'portrait photography of a woman, warm lighting, on the desert', 
+        1, 
+        20,
+        'images/restxt_4.webp',
+    ],
+    [
+        'images/original_3.jpg',
+        'uniform-lit', 
+        'portrait photography of a woman, at noon. bright sunlight, cityscape', 
+        1, 
+        20,
+        'images/restxt_5.webp',
+    ],
+    [
+        'images/original_3.jpg',
+        'normal', 
+        'portrait photography of a woman, at noon. bright sunlight, cityscape', 
+        1, 
+        20,
+        'images/restxt_6.webp',
+    ],
+    [
+        'images/original_4.jpg',
+        'normal', 
+        'portrait photography of a woman, at noon. bright sunlight, cityscape', 
+        1, 
+        20,
+        'images/restxt_7.webp',
+    ],
+    [
+        'images/original_4.jpg',
+        'normal', 
+        'portrait photography of a woman, neon lighting, by the city street', 
+        1, 
+        20,
+        'images/restxt_8.webp',
+    ],
+]
+
+image_guided_examples = [
+    [
+        'images/original_1.jpg',
+        'images/reference_2.jpg', 
+        12345, 
+        20,
+        'images/resimg_1.webp',
+    ],
+    [
+        'images/original_3.jpg',
+        'images/reference_4.jpg', 
+        12345, 
+        20,
+        'images/resimg_2.webp',
+    ],
+    [
+        'images/original_4.jpg',
+        'images/reference_3.jpg',
+        12345, 
+        20,
+        'images/resimg_3.webp',
+    ],
+    [
+        'images/original_2.jpg',
+        'images/reference_3.jpg', 
+        12345, 
+        20,
+        'images/resimg_4.webp',
+    ],
+]
+

ops.py

@@ -0,0 +1,135 @@
+import base64
+import numpy as np
+import cv2
+import json
+import pickle
+import itertools
+import requests
+from io import StringIO
+
+from io import BytesIO
+import re, os
+from PIL import Image
+from datetime import datetime
+import random
+
+import torch
+
+import socket
+r_b64_prefix = r'^data:image/.+;base64,'
+
+from PIL import ImageOps
+import oss2
+
+
+class OSS2:
+    def __init__(self, bucket):
+        self.bucket = bucket
+
+    def get_download(self, object_name, download_name=None, minute=5):
+        """
+        获取文件下载地址
+        :param object_name: OSS 路径
+        :param download_name:  下载重命名
+        :param minute:  请求超时时间
+        :return: 授权URL
+        """
+        params = None
+        if download_name:
+            params = {"response-content-disposition": 'attachment;filename="{0}"'.format(quote(download_name, 'utf-8'))}
+        url = self.bucket.sign_url('GET', object_name, minute * 60, params=params)
+        return url
+
+
+
+
+
+def _pil_image_push_ossdir(bucket, image_pil, object_key, quality=85, format='jpeg'):
+    import io
+    print_with_datetime('upload object_key {}'.format(object_key))
+    oss_path = None
+    for _ in range(5):
+        try:
+            img_byte_arr = io.BytesIO()
+            image_pil.save(img_byte_arr, format=format, quality=quality)  # 根据实际图片格式选择保存格式
+            img_byte_arr.seek(0)  # 回到开始位置
+            bucket.put_object(object_key, img_byte_arr)
+            oss_path = object_key
+            return oss_path
+        except Exception as e:
+            print_with_datetime('error, upload oss failed, {}'.format(str(e)))
+    return oss_path
+
+
+
+def resize_keep_hw_rate(image, tar_res):
+    w, h = image.size
+    if w > h:
+        tar_w = tar_res
+        tar_h = int(tar_w * h / w)
+    elif h > w:
+        tar_h = tar_res
+        tar_w = int(tar_h * w / h)
+    else:
+        tar_w = tar_res
+        tar_h = tar_res
+    image = image.resize((tar_w, tar_h), Image.LANCZOS)
+    return image
+
+
+def decode_img_from_url(url):
+    import requests
+    from io import BytesIO
+    for _ in range(5):
+        try:
+            content = requests.get(url, stream=True, timeout=10).content
+            img = Image.open(BytesIO(content))
+            img = ImageOps.exif_transpose(img)
+            break
+        except Exception as e:
+            print_with_datetime(f"url decode error {url} {str(e)}")
+            img = None
+    return img
+
+
+def make_sub_file_folder(file_path):
+    subdir = file_path.split(os.path.basename(file_path))[0]
+    os.makedirs(subdir, exist_ok=True)
+
+
+def get_random_file_name():
+    timing_cur = datetime.now().timestamp() * 1000000
+    timing_cur = str(int(timing_cur)) + f'-{random.randint(0,99999999)}'
+    return timing_cur
+
+
+def print_with_datetime(string):
+    print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}: {string}', flush=True)
+
+
+def b64_pil(base64Data, mode = 'RGB'):
+    try:
+        base64Data = re.sub(r_b64_prefix, '', base64Data)
+        imgData = base64.b64decode(base64Data)
+        image = Image.open(BytesIO(imgData))
+    except:
+        image = None
+    return image
+
+
+def pil_b64(pil_img, fmt='jpeg', quality=75):
+    output_buffer = BytesIO()
+    pil_img.save(output_buffer, format=fmt, quality=quality)
+    byte_data = output_buffer.getvalue()
+    base64_str = base64.b64encode(byte_data).decode('utf-8')
+    if fmt == 'jpeg':
+        fmt_out = 'jpg'
+    else:
+        fmt_out = fmt
+    return f'data:image/{fmt_out};base64,' + base64_str
+
+
+def get_ip():
+    hostname = socket.gethostname()
+    ip_address = socket.gethostbyname(hostname)
+    return ip_address

watermark.py

@@ -0,0 +1,49 @@
+import os
+import oss2
+from PIL import Image
+import gradio as gr
+import numpy as np
+from io import BytesIO
+
+
+
+class WatermarkApp:
+    def __init__(self):
+        pass
+    
+    def process_image(self, image):
+        watermark = Image.open('asset/chaojihui-v2.png')
+        # 获取水印的等比例缩放尺寸
+        wm_width, wm_height = watermark.size
+        target_size = 100
+        alpha = 0.3
+
+        if wm_width > wm_height:
+            new_wm_width = target_size
+            new_wm_height = int((wm_height / wm_width) * target_size)
+        else:
+            new_wm_height = target_size
+            new_wm_width = int((wm_width / wm_height) * target_size)
+
+        wm = watermark.resize((new_wm_width, new_wm_height), Image.Resampling.LANCZOS)
+
+        # 处理透明度
+        if alpha < 1.0:
+            wm = self.apply_transparency(wm, alpha)
+
+        # 确保原图是RGBA模式
+        if image.mode != 'RGBA':
+            image = image.convert('RGBA')
+
+        # 在整个图片上铺满水印
+        for y in range(0, image.height, new_wm_height):
+            for x in range(0, image.width, new_wm_width):
+                image.alpha_composite(wm, dest=(x, y))
+        return image
+
+    def apply_transparency(self, watermark, alpha):
+        """应用透明度"""
+        matrix = watermark.split()[-1].point(lambda x: x * alpha)
+        watermark.putalpha(matrix)
+        return watermark
+