Update app.py

app.py

@@ -347,74 +347,74 @@ def extract_features_siglip(image):
     return image_features
 
 @spaces.GPU
-def infer(crop_image_path,full_image_path,state,language,session_type,task_type=None):
+def infer(crop_image_path,full_image_path,state,language,task_type=None):
     print("task type",task_type)
     gallery_output = []
-    if session_type=="Session 1":
-        if task_type=="task 1":
-            gallery_output.append("recomendation_pic/1.8.jpg")
-            gallery_output.append("recomendation_pic/1.9.jpg")
-            input_image = Image.open(full_image_path).convert("RGB")
-            input_features = extract_features_siglip(input_image.convert("RGB"))
-            input_features = input_features.detach().cpu().numpy()
-            input_features = np.float32(input_features)
-            faiss.normalize_L2(input_features)
-            distances, indices = index.search(input_features, 2)      
-            for i,v in enumerate(indices[0]):
-                sim = -distances[0][i]
-                image_url = df.iloc[v]["Link"]
-                img_retrieved = read_image_from_url(image_url)
-                gallery_output.append(img_retrieved)
-            if language=="English":
-                msg="🖼️ Please refer to the section below to see the recommended results."
-            else:
-                msg="🖼️  请到下方查看推荐结果。"
-            state+=[(None,msg)]
-
-            return gallery_output,state,state
-        elif task_type=="task 2":
-            gallery_output.append("recomendation_pic/2.8.jpg")
-            gallery_output.append("recomendation_pic/2.9.png")
-            input_image = Image.open(full_image_path).convert("RGB")
-            input_features = extract_features_siglip(input_image.convert("RGB"))
-            input_features = input_features.detach().cpu().numpy()
-            input_features = np.float32(input_features)
-            faiss.normalize_L2(input_features)
-            distances, indices = index.search(input_features, 2)      
-            for i,v in enumerate(indices[0]):
-                sim = -distances[0][i]
-                image_url = df.iloc[v]["Link"]
-                img_retrieved = read_image_from_url(image_url)
-                gallery_output.append(img_retrieved)
-            if language=="English":
-                msg="🖼️ Please refer to the section below to see the recommended results."
-            else:
-                msg="🖼️  请到下方查看推荐结果。"
-            state+=[(None,msg)]
-
-            return gallery_output,state,state
-        
-        elif task_type=="task 3":
-            gallery_output.append("recomendation_pic/3.8.png")
-            gallery_output.append("recomendation_pic/3.9.png")
-            input_image = Image.open(full_image_path).convert("RGB")
-            input_features = extract_features_siglip(input_image.convert("RGB"))
-            input_features = input_features.detach().cpu().numpy()
-            input_features = np.float32(input_features)
-            faiss.normalize_L2(input_features)
-            distances, indices = index.search(input_features, 2)      
-            for i,v in enumerate(indices[0]):
-                sim = -distances[0][i]
-                image_url = df.iloc[v]["Link"]
-                img_retrieved = read_image_from_url(image_url)
-                gallery_output.append(img_retrieved)
-            if language=="English":
-                msg="🖼️ Please refer to the section below to see the recommended results."
-            else:
-                msg="🖼️  请到下方查看推荐结果。"
-            state+=[(None,msg)]
-
-            return gallery_output,state,state
+
+    if task_type=="task 1":
+        gallery_output.append("recomendation_pic/1.8.jpg")
+        gallery_output.append("recomendation_pic/1.9.jpg")
+        input_image = Image.open(full_image_path).convert("RGB")
+        input_features = extract_features_siglip(input_image.convert("RGB"))
+        input_features = input_features.detach().cpu().numpy()
+        input_features = np.float32(input_features)
+        faiss.normalize_L2(input_features)
+        distances, indices = index.search(input_features, 2)      
+        for i,v in enumerate(indices[0]):
+            sim = -distances[0][i]
+            image_url = df.iloc[v]["Link"]
+            img_retrieved = read_image_from_url(image_url)
+            gallery_output.append(img_retrieved)
+        if language=="English":
+            msg="🖼️ Please refer to the section below to see the recommended results."
+        else:
+            msg="🖼️  请到下方查看推荐结果。"
+        state+=[(None,msg)]
+
+        return gallery_output,state,state
+    elif task_type=="task 2":
+        gallery_output.append("recomendation_pic/2.8.jpg")
+        gallery_output.append("recomendation_pic/2.9.png")
+        input_image = Image.open(full_image_path).convert("RGB")
+        input_features = extract_features_siglip(input_image.convert("RGB"))
+        input_features = input_features.detach().cpu().numpy()
+        input_features = np.float32(input_features)
+        faiss.normalize_L2(input_features)
+        distances, indices = index.search(input_features, 2)      
+        for i,v in enumerate(indices[0]):
+            sim = -distances[0][i]
+            image_url = df.iloc[v]["Link"]
+            img_retrieved = read_image_from_url(image_url)
+            gallery_output.append(img_retrieved)
+        if language=="English":
+            msg="🖼️ Please refer to the section below to see the recommended results."
+        else:
+            msg="🖼️  请到下方查看推荐结果。"
+        state+=[(None,msg)]
+
+        return gallery_output,state,state
+    
+    elif task_type=="task 3":
+        gallery_output.append("recomendation_pic/3.8.png")
+        gallery_output.append("recomendation_pic/3.9.png")
+        input_image = Image.open(full_image_path).convert("RGB")
+        input_features = extract_features_siglip(input_image.convert("RGB"))
+        input_features = input_features.detach().cpu().numpy()
+        input_features = np.float32(input_features)
+        faiss.normalize_L2(input_features)
+        distances, indices = index.search(input_features, 2)      
+        for i,v in enumerate(indices[0]):
+            sim = -distances[0][i]
+            image_url = df.iloc[v]["Link"]
+            img_retrieved = read_image_from_url(image_url)
+            gallery_output.append(img_retrieved)
+        if language=="English":
+            msg="🖼️ Please refer to the section below to see the recommended results."
+        else:
+            msg="🖼️  请到下方查看推荐结果。"
+        state+=[(None,msg)]
+
+        return gallery_output,state,state
     
     elif crop_image_path:
         input_image = Image.open(crop_image_path).convert("RGB")
@@ -1090,7 +1090,7 @@ def inference_click(image_input, point_prompt, click_mode, enable_wiki, language
         Image.open(out["crop_save_path"]).save(new_crop_save_path)
         print("new crop save",new_crop_save_path)
 
-    yield state, state, click_state, image_input_nobackground, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state,new_crop_save_path,image_input_nobackground
+    return state, state, click_state, image_input_nobackground, click_index_state, input_mask_state, input_points_state, input_labels_state, out_state,new_crop_save_path,image_input_nobackground
     
 
 query_focus_en = [
@@ -1646,7 +1646,7 @@ async def texttospeech(text, language,gender='female'):
         return None
 
 # give the reason of recommendation
-async def associate(image_path,new_crop,openai_api_key,language,autoplay,length,log_state,sort_score,narritive,evt: gr.SelectData):
+async def associate(image_path,new_crop,openai_api_key,language,autoplay,length,log_state,sort_score,narritive,state,evt: gr.SelectData):
     persona=naritive_mapping[narritive]
     rec_path=evt._data['value']['image']['path']
     index=evt.index
@@ -1658,7 +1658,7 @@ async def associate(image_path,new_crop,openai_api_key,language,autoplay,length,
         image_paths=[image_path,rec_path]
     result=get_gpt_response(openai_api_key, image_paths, prompt)
     print("recommend result",result)
-    reason = [(None, f"{result}")]
+    state += [(None, f"{result}")]
     log_state = log_state + [(narritive, None)]
     log_state = log_state + [(f"image sort ranking {sort_score}", None)]
     log_state = log_state + [(None, f"{result}")]
@@ -1668,11 +1668,11 @@ async def associate(image_path,new_crop,openai_api_key,language,autoplay,length,
     audio_output=None
     if autoplay:
         audio_output = await texttospeech(read_info, language)
-    return reason,audio_output,log_state,index,gr.update(value=[])
+    return state,state,audio_output,log_state,index,gr.update(value=[])
 
-def change_naritive(session_type,image_input, chatbot, state, click_state, paragraph, origin_image,narritive,language="English"):
+def change_naritive(session_type,image_input, state, click_state, paragraph, origin_image,narritive,task_instruct,gallery_output,reco_reasons,language="English"):
     if session_type=="Session 1":
-        return None, [], [], [[], [], []], "", None, []
+        return None, [], [], [[], [], []], "", None, None, [], [],[]
     else:
         if language=="English":
             if narritive=="Third-person" :
@@ -1720,7 +1720,7 @@ def change_naritive(session_type,image_input, chatbot, state, click_state, parag
             ]
 
 
-        return image_input, state, state, click_state, paragraph, origin_image
+        return image_input, state, state, click_state, paragraph, origin_image,task_instruct,gallery_output,reco_reasons,reco_reasons
 
 
 def print_like_dislike(x: gr.LikeData,state,log_state):
@@ -1766,7 +1766,7 @@ def create_ui():
     examples = [
         ["test_images/1.The Ambassadors.jpg","test_images/task1.jpg","task 1"],
         ["test_images/2.Football Players.jpg","test_images/task2.jpg","task 2"],
-        ["test_images/3.Along the River during the Qingming Festival.jpeg","test_images/task3.jpg","task 3"],
+        ["test_images/3-square.jpg","test_images/task3.jpg","task 3"],
         # ["test_images/test4.jpg"],
         # ["test_images/test5.jpg"],
         # ["test_images/Picture5.png"],
@@ -1810,6 +1810,7 @@ def create_ui():
         # store the whole image path 
         image_path=gr.State('')
         pic_index=gr.State(None)
+        recomended_state=gr.State([])
    
         
         with gr.Row():
@@ -1821,8 +1822,7 @@ def create_ui():
             )
         with gr.Row():  
             with gr.Column(scale=1,min_width=50,visible=False) as instruct:
-                task_instuction=gr.Image(type="pil", interactive=True, elem_classes="task_instruct",height=650,label=None)
-                 
+                task_instuction=gr.Image(type="pil", interactive=True, elem_classes="task_instruct",height=650,label=None)                    
             with gr.Column(scale=6):
                 with gr.Column(visible=False) as modules_not_need_gpt:
                     
@@ -1941,6 +1941,7 @@ def create_ui():
             with gr.Column(scale=4):                    
                 with gr.Column(visible=True) as module_key_input:
                     openai_api_key = gr.Textbox(
+                        value="sk-proj-bxHhgjZV8TVgd1IupZrUT3BlbkFJvrthq6zIxpZVk3vwsvJ9",
                         placeholder="Input openAI API key",
                         show_label=False,
                         label="OpenAI API Key",
@@ -2206,14 +2207,14 @@ def create_ui():
         # )
         recommend_btn.click(
             fn=infer,
-            inputs=[new_crop_save_path,image_path,state,language,session_type,task_type],
+            inputs=[new_crop_save_path,image_path,state,language,task_type],
             outputs=[gallery_result,chatbot,state]
             )
 
         gallery_result.select(
             associate,
-            inputs=[image_path,new_crop_save_path,openai_api_key,language,auto_play,length,log_state,sort_rec,naritive],
-            outputs=[recommend_bot,output_audio,log_state,pic_index,recommend_score],
+            inputs=[image_path,new_crop_save_path,openai_api_key,language,auto_play,length,log_state,sort_rec,naritive,recomended_state],
+            outputs=[recommend_bot,recomended_state,output_audio,log_state,pic_index,recommend_score],
             
             
         )
@@ -2434,11 +2435,18 @@ def create_ui():
 
         # cap_everything_button.click(cap_everything, [paragraph, visual_chatgpt, language,auto_play], 
         #                             [paragraph_output,output_audio])
-        
+        def reset_and_add(origin_image):
+            new_prompt = "Positive"
+            new_add_icon = "assets/icons/plus-square-blue.png"
+            new_add_css = "tools_button_clicked"
+            new_minus_icon = "assets/icons/minus-square.png"
+            new_minus_css= "tools_button" 
+            return [[],[],[]],origin_image, new_prompt, gr.update(icon=new_add_icon,elem_classes=new_add_css), gr.update(icon=new_minus_icon,elem_classes=new_minus_css)
+            
         clear_button_click.click(
-            lambda x: ([[], [], []], x),
+            reset_and_add,
             [origin_image],
-            [click_state, image_input],
+            [click_state, image_input,point_prompt,add_button,minus_button],
             queue=False,
             show_progress=False
         )
@@ -2525,10 +2533,10 @@ def create_ui():
                             paragraph,artist,gender,image_path, log_state,history_log,output_audio])
 
         example_image.change(clear_chat_memory, inputs=[visual_chatgpt])
-        example_image.change(
-            lambda:([],[]), 
-            [],
-            [gallery_result,recommend_bot])
+        # example_image.change(
+        #     lambda:([],[]), 
+        #     [],
+        #     [gallery_result,recommend_bot])
 
         # def on_click_tab_selected():
         #     if gpt_state ==1:
@@ -2672,20 +2680,21 @@ def create_ui():
         
         naritive.change(
             change_naritive,
-            [session_type, image_input, chatbot, state, click_state, paragraph, origin_image,naritive,language],
-            [image_input, chatbot, state, click_state, paragraph, origin_image,gallery_result],
+            [session_type, image_input, state, click_state, paragraph, origin_image,naritive,
+             task_instuction,gallery_result,recomended_state,language],
+            [image_input, chatbot, state, click_state, paragraph, origin_image,task_instuction,gallery_result,recomended_state,recommend_bot],
             queue=False,
             show_progress=False
         
         )
         def session_change():
             instruction=Image.open('test_images/task4.jpg')
-            return None, [], [], [[], [], []], "", None, [],[],instruction
+            return None, [], [], [[], [], []], "", None, [],[],instruction,"task 4"
         
         session_type.change(
             session_change,
             [],
-            [image_input, chatbot, state, click_state, paragraph, origin_image,history_log,log_state,task_instuction]
+            [image_input, chatbot, state, click_state, paragraph, origin_image,history_log,log_state,task_instuction,task_type]
         )
         
         # upvote_btn.click(