Update app.py

app.py

@@ -42,23 +42,29 @@ COLOURS_DICT = color_data['color_mapping']
 ATTRIBUTES_DICT = attributes_data['attribute_mapping']
 
 
-def shot(input, category):
-    subColour,mainColour,score = get_colour(ast.literal_eval(str(input)),category)
-    common_result = get_predicted_attributes(ast.literal_eval(str(input)),category)
-    openai_parsed_response = get_openAI_tags(ast.literal_eval(str(input)))
-    face_embeddings = get_face_embeddings(ast.literal_eval(str(input)))
-    cropped_images = get_cropped_images(ast.literal_eval(str(input)),category)
-    print(cropped_images)
-    return {
-        "colors":{
+def shot(input, category, level):
+    output_dict = {}
+    if level = 'variant':
+        subColour,mainColour,score = get_colour(ast.literal_eval(str(input)),category)    
+        openai_parsed_response = get_openAI_tags(ast.literal_eval(str(input)))
+        face_embeddings = get_face_embeddings(ast.literal_eval(str(input)))
+        cropped_images = get_cropped_images(ast.literal_eval(str(input)),category)
+        output_dict['colors'] = {
             "main":mainColour,
             "sub":subColour,
-            "score":round(score*100,2)
-        }, 
-        "attributes":common_result,
-        "image_mapping":openai_parsed_response,
-        "face_embeddings":face_embeddings
-    }
+            "score":score
+        }
+        output_dict['image_mapping'] = openai_parsed_response
+        output_dict['face_embeddings'] = face_embeddings
+        output_dict['cropped_images'] = cropped_images
+
+
+    if level = 'product':
+        common_result = get_predicted_attributes(ast.literal_eval(str(input)),category)    
+        output_dict['attributes'] = common_result
+        output_dict['subcategory'] = category
+    
+    return output_dict
 
 
 
@@ -435,6 +441,7 @@ iface = gr.Interface(
     inputs=[
         gr.Textbox(label="Image URLs (starting with http/https) comma seperated "), 
         gr.Textbox(label="Category")
+        gr.Textbox(label="Level; accepted 'variant' or 'product'")
     ], 
     outputs="text" ,
      examples=[