ui.py

import gradio as gr
from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
from PIL import Image
import torch
import torch.nn.functional as F
from torchvision.io import read_image, ImageReadMode
from torchvision.transforms.functional import to_pil_image
from segment_anything import sam_model_registry, SamPredictor
from ibm_watsonx_ai import Credentials
from ibm_watsonx_ai.foundation_models import ModelInference
from ibm_watsonx_ai.foundation_models.schema import TextChatParameters
from openai import OpenAI

import sys
sys.path.append("..")
import numpy as np
import json
import requests
import os
import heapq
import subprocess
from io import BytesIO
import base64

OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
project_id = "skills-network"


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load CLIPSeg Model
processor_clipseg = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
model_clipseg = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")

# Load SAM
# The URL of the SAM model parameters
url = "https://cf-courses-data.s3.us.cloud-object-storage.appdomain.cloud/ac2ryJJQCwFcauPZLw9DRg/sam-vit-b-01ec64.pth"
# The desired output filename (optional, wget will infer if not provided)
file_name = "sam-vit-b-01ec64.pth"

if not os.path.exists(file_name):
    subprocess.run(["wget", url])

sam_checkpoint = "sam-vit-b-01ec64.pth"
model_type = "vit_b"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

# Load LLM
credentials = Credentials(
                url = "https://us-south.ml.cloud.ibm.com"
                )

# Get sample parameter values
sample_params = TextChatParameters.get_sample_params()

# Initialize the TextChatParameters object with the sample values
params = TextChatParameters(**sample_params)

model_chat = ModelInference(
    model_id='ibm/granite-4-h-small',  #meta-llama/llama-4-maverick-17b-128e-instruct-fp8
    credentials=credentials,
    project_id=project_id,
    params=params,
)

def process_gpt_images(response, filename, image_dir):
    os.makedirs(image_dir, exist_ok=True)  # ensure the directory exists

    images = [base64.b64decode(datum.b64_json) for datum in response.data]  # decode base64 images
    image_names = [f"{filename}_{i + 1}.png" for i in range(len(images))]  # create names
    filepaths = [os.path.join(image_dir, name) for name in image_names]  # create filepaths
    for image, filepath in zip(images, filepaths):  # loop through the images
        with open(filepath, "wb") as image_file:  # open the file
            image_file.write(image)  # write the image to the file
    return filepaths

def extract_mask_only(masked_image):
    """
    Extracts the mask (alpha channel) from the PIL Image object 
    returned by gr.ImageMask.
    """
    if masked_image is None:
        return np.zeros((100, 100), dtype=np.uint8) # Return a blank mask if no image
    
    # 1. Convert the image to RGBA (Red, Green, Blue, Alpha)
    # The mask is stored in the Alpha channel.
    masked_image = Image.fromarray(masked_image)
    img_rgba = masked_image.convert("RGBA")
    
    # 2. Convert to a NumPy array
    img_array = np.array(img_rgba)
    
    # 3. The Alpha channel is the 4th channel (index 3)
    # This channel's values will be 0 where there is no mask, and > 0 where there is a mask.
    mask_channel = img_array[:, :, 3] 

    # # If you want a standard 8-bit grayscale image of the mask:
    # mask_image = Image.fromarray(mask_channel)

    return mask_channel


def fashion_design(query, mask_input, image_path):
    #### Understand user's query
    error_flag = False
    prompt = f"""
    I give you a user query: {query}
    Fill in the blanks.
    Can you point to me (CONCISE): 
    (1) What needs to be changed: [BLANK]; 
    (2) Will be changed to what: [BLANK].

    Structure your answers to the blanks as a Python Dictionary, with keys '1' and '2'.
    """

    response = model_chat.chat(
    messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                ],
            }
        ]
    )
    response_output = response['choices'][0]['message']['content']
    try:
        response_output = json.loads(response_output)
    except:
        error_warning = "Please format your query as: change A to B."
        error_flag = True
    
    if error_flag:
        return None, error_warning
    
    keyerr_flag = False
    try:
        to_be_changed = response_output["1"]
        will_change_to = response_output["2"]
    except:
        keyerr_flag = True
    
    if keyerr_flag:
        keys = list(response_output.keys())
        to_be_changed = response_output[keys[0]]
        will_change_to = response_output[keys[1]]

    if len(mask_input['layers']) > 0:
        num_masks = 2
        chosen_mask = extract_mask_only(mask_input['layers'][0])
        chosen_mask = chosen_mask.astype("uint8")
        chosen_mask[chosen_mask != 0] = 255
        chosen_mask[chosen_mask == 0] = 1
        chosen_mask[chosen_mask == 255] = 0
        chosen_mask[chosen_mask == 1] = 255

        # create a base blank mask
        width = chosen_mask.shape[1]
        height = chosen_mask.shape[0]
        mask = Image.new("RGBA", (width, height), (0, 0, 0, 255))  # create an opaque image mask

        # Convert mask back to pixels to add our mask replacing the third dimension
        pix = np.array(mask)
        pix[:, :, 3] = chosen_mask

        new_mask = Image.fromarray(pix)
        new_mask.save(f"mask_candidates/mask_option{0}.png")
    else:
        #### CLIPSeg
        # Read as a uint8 tensor [C, H, W] in RGB (forces 3 channels; drops alpha if present)
        img_tensor = read_image(image_path)

        # Convert to PIL.Image in RGB
        image_pil = to_pil_image(img_tensor)

        # Send the text_prompt and the image_pil to the model for prediction
        inputs = processor_clipseg(text=[to_be_changed], images=[image_pil], return_tensors="pt")
        with torch.no_grad():
            outputs = model_clipseg(**inputs)
            preds = torch.sigmoid(outputs.logits)  # mask probabilities

        # Convert the predicted mask to be a binary mask of only 0 or 1.
        mask = preds.squeeze().cpu().numpy()
        mask_binary = (mask > 0.05).astype(np.uint8)  # threshold mask

        # Generate the mask from the mask_binary
        mask_pil = Image.fromarray(mask_binary*255)
        rough_mask_np = np.array(mask_pil.convert("L")) > 128

        #### SAM
        predictor = SamPredictor(sam)
        numpy_img = torch.permute(img_tensor,[1,2,0]).numpy()
        predictor.set_image(numpy_img)

        w, h, _ = numpy_img.shape
        wm, hm = rough_mask_np.shape

        # Derive bounding box from rough mask
        y, x = np.where(rough_mask_np)
        bbox = np.array([min(x)/hm*h, min(y)/wm*w, max(x)/hm*h, max(y)/wm*w])  # [x0, y0, x1, y1]

        masks, _, _ = predictor.predict(
            point_coords=None,
            point_labels=None,
            box=bbox[None, :],
            multimask_output=True,
        )

        num_masks = len(masks)
        for i in range(num_masks):
            # We'll now reverse the mask so that it is clear and everything else is white
            chosen_mask = masks[i]
            chosen_mask = chosen_mask.astype("uint8")
            chosen_mask[chosen_mask != 0] = 255
            chosen_mask[chosen_mask == 0] = 1
            chosen_mask[chosen_mask == 255] = 0
            chosen_mask[chosen_mask == 1] = 255

            # create a base blank mask
            width = chosen_mask.shape[1]
            height = chosen_mask.shape[0]
            mask = Image.new("RGBA", (width, height), (0, 0, 0, 1))  # create an opaque image mask

            # Convert mask back to pixels to add our mask replacing the third dimension
            pix = np.array(mask)
            pix[:, :, 3] = chosen_mask

            # Convert pixels back to an RGBA image and display
            new_mask = Image.fromarray(pix)
            new_mask.save(f"mask_candidates/mask_option{i}.png")


    client = OpenAI(api_key=OPENAI_API_KEY)
    finalized = []
    n = 3
    ERROR_FLAG = False
    print('Editing....')
    for i in range(num_masks-1):
        mask_path = f"mask_candidates/mask_option{i}.png"
        try:
            response = client.images.edit(
                model='gpt-image-2',
                image=open(image_path, "rb"),  # from the generation section
                mask=open(mask_path, "rb"),  # from right above
                prompt=will_change_to,  # provide a prompt to fill the space
                n=n,
                size="1024x1024",
            )
            image_paths = process_gpt_images(response, f"edits_mask{i}", 'edited_images')
            print(image_paths)

            for edited_path in image_paths:
                    finalized.append(edited_path)
        except Exception as e:
            print(f"API error on mask {i}: {e}") 
            ERROR_FLAG = True
            break

    print('Out....')
    if len(finalized) == 0:
        message = 'Edited, but none of the generated images satisfy your request... Please try again by changing your prompt.'
    elif ERROR_FLAG:
        message = "Don't worry! Please TRY again, some minor issue with the connection..."
    else:
        message = 'Edited Sucessfully! (Please note that the models are imperfect, so the image quality can vary significantly)'

    return finalized, message


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## Fashion Design Assistant")
    gr.Markdown("#### IMPORTANT: Make sure your uploaded image is less than 4 MB!")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload an Image", type="filepath")
            gr.Markdown('#### Optional: If you want to draw your own mask for accuracy, please draw here. You can directly use the paint brush embedded in the window below and paint the region for the mask. You may customize the size of the brush.')
            mask_input = gr.ImageMask(label="Draw mask on image")
            query_input = gr.Textbox(label="Enter your prompt", placeholder="e.g., change the shirt to a sweater")
            submit_button = gr.Button("DESIGN!", variant="primary")

        with gr.Column():
            gr.Markdown("### AI Fashion Gallery")
            gr.Markdown("#### Each editing could take about 120 seconds (2 minutes).")
            gallery_output = gr.Gallery(
                label="AI Designs",
                show_label=False,
                elem_id="gallery",
                columns=[3],
                rows=[1],
                object_fit="contain",
                height="auto"
            )
            message_output = gr.Markdown(label="Status Message")

    # When button is clicked, call function f
    submit_button.click(
        fn=fashion_design,
        inputs=[query_input,mask_input,image_input],
        outputs=[gallery_output,message_output]
    )

    # 🔹 Example (image, query) pairs
    examples = [
        ["example_images/cat-blue-cloth.png", "change the blue cloth to a shirt with pumpkins"],
        ["example_images/lady-white-shirt.png", "change the white shirt to a pink sweater"]
    ]

    # Built-in Gradio Examples block
    gr.Examples(
        examples=examples,
        inputs=[image_input, query_input],
        label="Try These Example Image–Prompt Pairs"
    )


# --------------------------------------------------
# Launch the app
# --------------------------------------------------
if __name__ == "__main__":
    demo.launch(share=True)