Update app.py

app.py

@@ -1,41 +1,69 @@
-from turtle import title
+import json
+import spaces
 import requests
-from io import BytesIO
-import gradio as gr
-from transformers import pipeline
 import numpy as np
+import gradio as gr
 from PIL import Image
-import spaces
-
+from io import BytesIO
+from turtle import title
+from transformers import pipeline
+import ast
 pipe = pipeline("zero-shot-image-classification", model="patrickjohncyh/fashion-clip")
-images="dog.jpg"
+
+file_path = 'config.json'
+
+# Open and read the JSON file
+with open(file_path, 'r') as file:
+    data = json.load(file)
+
+COLOURS_DICT = data['color_mapping']
+
+
+
+def shot(input, category):
+    subColour,mainColour,score = get_colour(ast.literal_eval(str(input)),category)
+    return subColour,mainColour,score
+
+
 
 @spaces.GPU  
-def shot(input, labels_text):
-    if isinstance(input, str) and (input.startswith("http://") or input.startswith("https://")):
-        # Input is a URL
-        response = requests.get(input)
-        PIL_image = Image.open(BytesIO(response.content)).convert('RGB')
-    else:
-        # Input is an uploaded image
-        PIL_image = Image.fromarray(np.uint8(input)).convert('RGB')
-    
-    labels = labels_text.split(",")
-    res = pipe(images=PIL_image, 
-               candidate_labels=labels,
-               hypothesis_template="This is a photo of a {}")
-    return {dic["label"]: dic["score"] for dic in res}
+def get_colour(image_urls, category):
+    colourLabels = list(COLOURS_DICT.keys())
+    for i in range(len(colourLabels)):
+        colourLabels[i] = colourLabels[i] + " clothing: " + category
+
+    responses = pipe(image_urls, candidate_labels=colourLabels, device=device)
+    # Get the most common colour
+    mainColour = responses[0][0]['label'].split(" clothing:")[0]
+
+
+    if mainColour not in COLOURS_DICT:
+        return None, None, None
+
+    # Add category to the end of each label
+    labels = COLOURS_DICT[mainColour]
+    for i in range(len(labels)):
+        labels[i] = labels[i] + " clothing: " + category
+
+    # Run pipeline in one go
+    responses = pipe(image_urls, candidate_labels=labels, device=device)
+    subColour = responses[0][0]['label'].split(" clothing:")[0]
+
+    return subColour, mainColour, responses[0][0]['score']
+
+
+
 
 # Define the Gradio interface with the updated components
 iface = gr.Interface(
     fn=shot, 
     inputs=[
-        gr.Textbox(label="Image URL (starting with http/https) or Upload Image"), 
-        gr.Textbox(label="Labels (comma-separated)")
+        gr.Textbox(label="Image URLs (starting with http/https) comma seperated "), 
+        gr.Textbox(label="Category")
     ], 
     outputs=gr.Label(), 
     description="Add an image URL (starting with http/https) or upload a picture, and provide a list of labels separated by commas.",
-    title="Zero-shot Image Classification"
+    title="Full product flow"
 )
 
 # Launch the interface