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)