app.py

import gradio as gr
import json
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
from PIL import Image
import torch
import numpy as np
import requests
import cv2
import urllib.request
import os

# -----------------------------------------------------------------------------
# 1. Define Functions

# Load the catalog from the JSON file
def load_catalog(filepath="catalog.json"):
    """Loads the product catalog from a JSON file."""
    with open(filepath, 'r') as f:
        catalog = json.load(f)
    return catalog

# LLM Component (Using Flan-T5 - CPU Friendly)
model_name = "google/flan-t5-small"  # CPU-friendly model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name) # No device_map needed

llm_pipeline = pipeline(
    "text2text-generation",  # Use text2text-generation for Flan-T5
    model=model,
    tokenizer=tokenizer,
    #device=0, # Remove device specification - CPU only
    truncation=True,
    max_length=200
)

def get_tile_recommendations(query, catalog):
    """Uses an LLM to find tiles that match a text query."""
    query = query.lower()
    tile_ids = []
    # Keyword Algorithm, since the LLM is unreliable.
    for tile in catalog:
        if query.lower() in tile['tile_id'].lower(): #Test it
            tile_ids.append(tile["tile_id"])

    return tile_ids

def get_llm_response(query):
    """Simply ask the LLM the question"""
    prompt = f"Answer this question {query}"
    response = llm_pipeline(prompt, max_length=200)
    llm_response = response[0]['generated_text']
    return llm_response

def detect_floor_color(image):
    """Detects the floor region based on color in the lower part of the image."""
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    lower_color = np.array([0, 0, 0])  # Adjust these values (H,S,V)
    upper_color = np.array([180, 50, 200])  # Adjust these values
    mask = cv2.inRange(hsv_image, lower_color, upper_color)
        # Apply morphological operations to clean the mask
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.erode(mask, kernel, iterations=2)
    mask = cv2.dilate(mask, kernel, iterations=2)

    # Find contours
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # Filter contours based on location (lower part of image) and size
    image_height = image.shape[0]
    floor_contours = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if y > image_height // 2 and cv2.contourArea(contour) > 1000:  #Adjust min Area
            floor_contours.append(contour)
    if floor_contours:
            # Select the contour with the largest area
            floor_contour = max(floor_contours, key=cv2.contourArea)
            x, y, w, h = cv2.boundingRect(floor_contour)
            return y, h
    else:
        return None, None  #No floor found

def visualize_tile_on_floor(room_image, tile_image_url, floor_top=None, floor_height=None):
    """Overlays a tile image onto the floor of a room image, using detected floor if floor values are None."""
    try:

        room_image = cv2.cvtColor(np.array(room_image), cv2.COLOR_RGB2BGR)

        if floor_top is None or floor_height is None:
             floor_top, floor_height = detect_floor_color(room_image) #use floor detection
             if floor_top is None or floor_height is None: #If floor detection fails
                print ("Floor detection failed")
                return None

        # Load image via URL
        req = urllib.request.urlopen(tile_image_url)
        arr = np.asarray(bytearray(req.read()), dtype=np.uint8)
        tile_image = cv2.imdecode(arr, -1)
        tile_image = cv2.cvtColor(tile_image, cv2.COLOR_BGR2RGB)

        # Define the floor region
        floor_x1 = 0  # Start at the left edge of the image
        floor_x2 = room_image.shape[1]  # Go to the right edge of the image
        floor_y1 = floor_top
        floor_y2 = floor_top + floor_height

        # Extract the floor region
        floor_roi = room_image[floor_y1:floor_y2, floor_x1:floor_x2]

        # Ensure the floor ROI is not empty
        if floor_roi.size == 0:
            print("Error: Empty floor region selected.")
            return None

        # Resize tile to fit floor
        tile_image = cv2.resize(tile_image, (floor_x2-floor_x1, floor_y2-floor_y1))

        # Overlay the tile image onto the floor region
        try:
            room_image[floor_y1:floor_y2, floor_x1:floor_x2] = tile_image
        except ValueError as e:
            print(f"Value Error{e}")

        room_image = cv2.cvtColor(room_image, cv2.COLOR_BGR2RGB)  # back to cv2
        return room_image

    except Exception as e:
        print(f"Error in visualization: {e}")
        return None

# -----------------------------------------------------------------------------
# 2. Main Function
def predict(room_image, query):
    """Main function to orchestrate the AI sales assistant."""
    try:
        recommended_tile_ids = get_tile_recommendations(query, catalog)

        if recommended_tile_ids:
            first_tile_id = recommended_tile_ids[0]
            tile_info = next((tile for tile in catalog if tile["tile_id"] == first_tile_id), None)
            if tile_info:
                 # Now, use the image_url value from the catalog, instead of passing it as a separate argument
                tile_image_url = tile_info["image_url"] # get the image URL

                print (f"Found {tile_image_url}") #to debug
                visualization = visualize_tile_on_floor(room_image, tile_image_url) #floor detection
                if visualization is None:
                    return "Visualization failed.", None, "Error"
            else:
                return "Tile not found", None, "Error"
        else:
            return "No tiles matched query.", None, "No matches"

        llm_response = get_llm_response(query)

        return llm_response, visualization, "Success!"

    except Exception as e:
        return str(e), None, "Error"

# -----------------------------------------------------------------------------
# 3. Load resources
catalog = load_catalog()

# -----------------------------------------------------------------------------
# 4. Gradio Interface

iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload a Room Image"),
        gr.Textbox(lines=2, label="Describe the tiles you're looking for"),
    ],
    outputs=[
        gr.Textbox(label="Chatbot Response"),
        gr.Image(label="Visualization"),
        gr.Textbox(label="Status Message")]
    ,
    title="AI Tile Sales Assistant (Simplified)",
    description="Upload a room image and describe the tiles you want.  If no floor is automatically detected, adjust the HSV values."
)

if __name__ == "__main__":
    iface.launch(debug=True)